fix the bug from metis if num_parts == 1

Author: james77777778

test/test_metis.py

@@ -30,3 +30,8 @@ def test_metis(device):
                                          weighted=False)
         assert partptr.numel() == 3
         assert perm.numel() == 6
+
+        _, partptr, perm = mat.partition(num_parts=1, recursive=False,
+                                         weighted=True)
+        assert partptr.numel() == 2
+        assert perm.numel() == 6

torch_sparse/metis.py

@@ -33,8 +33,13 @@ def partition(src: SparseTensor, num_parts: int, recursive: bool = False,
     else:
         value = None
 
-    cluster = torch.ops.torch_sparse.partition(rowptr, col, value, num_parts,
-                                               recursive)
+    if num_parts > 1:
+        cluster = torch.ops.torch_sparse.partition(rowptr, col, value,
+                                                   num_parts, recursive)
+    elif num_parts == 1:
+        cluster = torch.zeros((src.size(0)), dtype=torch.long)
+    else:
+        raise ValueError
     cluster = cluster.to(src.device())
 
     cluster, perm = cluster.sort()