fix convert_partition.py

dmlc · Oct 17, 2024 · 7f339b9 · 7f339b9
1 parent 97d000b
commit 7f339b9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tools/distpartitioning/convert_partition.py b/tools/distpartitioning/convert_partition.py
@@ -351,6 +351,8 @@ def _process_partition_gb(
         sorted_idx = (
             th.repeat_interleave(indptr[:-1], split_size, dim=0) + sorted_idx
         )
+    else:
+        sorted_idx = th.arange(len(edge_ids))
 
     return indptr, indices[sorted_idx], edge_ids[sorted_idx]