fix

torch_sparse/__init__.py

@@ -55,7 +55,7 @@ from .mul import mul, mul_, mul_nnz, mul_nnz_  # noqa
 from .reduce import sum, mean, min, max  # noqa
 from .matmul import matmul  # noqa
 from .cat import cat, cat_diag  # noqa
-from .metis import partition_kway  # noqa
+from .metis import partition  # noqa
 
 from .convert import to_torch_sparse, from_torch_sparse  # noqa
 from .convert import to_scipy, from_scipy  # noqa
@@ -95,7 +95,7 @@ __all__ = [
     'matmul',
     'cat',
     'cat_diag',
-    'partition_kway',
+    'partition',
     'to_torch_sparse',
     'from_torch_sparse',
     'to_scipy',