[GraphBolt] Fix MiniBatch blocks. (#6766)

Co-authored-by: Ubuntu <[email protected]>
dmlc · Dec 18, 2023 · 1db71cf · 1db71cf
1 parent 43fb73d
commit 1db71cf
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/python/dgl/graphbolt/minibatch.py b/python/dgl/graphbolt/minibatch.py
@@ -395,7 +395,12 @@ def blocks(self):
                             (
                                 v.indptr,
                                 v.indices,
-                                torch.tensor([]),
+                                torch.arange(
+                                    0,
+                                    v.indptr[-1],
+                                    device=v.indptr.device,
+                                    dtype=v.indptr.dtype,
+                                ),
                             ),
                         )
                         for etype, v in subgraph.node_pairs.items()
@@ -421,7 +426,12 @@ def blocks(self):
                         (
                             node_pairs.indptr,
                             node_pairs.indices,
-                            torch.tensor([]),
+                            torch.arange(
+                                0,
+                                node_pairs.indptr[-1],
+                                device=node_pairs.indptr.device,
+                                dtype=node_pairs.indptr.dtype,
+                            ),
                         ),
                     )
                 num_src_nodes = original_row_node_ids.size(0)