fix test

test/test_metis.py

@@ -8,6 +8,10 @@ from .utils import devices
 @pytest.mark.parametrize('device', devices)
 def test_metis(device):
     mat = SparseTensor.from_dense(torch.randn((6, 6), device=device))
-    mat, partptr, perm = mat.partition_kway(num_parts=2)
+    mat, partptr, perm = mat.partition(num_parts=2, recursive=False)
+    assert partptr.numel() == 3
+    assert perm.numel() == 6
+
+    mat, partptr, perm = mat.partition(num_parts=2, recursive=True)
     assert partptr.numel() == 3
     assert perm.numel() == 6

torch_sparse/metis.py

@@ -21,4 +21,5 @@ def partition(
     return out, partptr, perm
 
 
-SparseTensor.partition = lambda self, num_parts: partition(self, num_parts)
+SparseTensor.partition = lambda self, num_parts, recursive=False: partition(
+    self, num_parts, recursive)