Fix shape of node_weight in METIS (#342)

* update * update
rusty1s · Sep 12, 2023 · 3246352 · 3246352
1 parent 14781cb
commit 3246352
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
@@ -48,6 +48,7 @@ jobs:
 
       - name: Install main package
         run: |
+          pip install scipy
           python setup.py develop
         env:
           WITH_METIS: 1

diff --git a/torch_sparse/metis.py b/torch_sparse/metis.py
@@ -52,7 +52,7 @@ def partition(
         value = None
 
     if balance_edge:
-        node_weight = col.new_zeros(col.size(0))
+        node_weight = col.new_zeros(rowptr.numel() - 1)
         node_weight.scatter_add_(0, col, torch.ones_like(col))
 
     if node_weight is not None: