Update @withMETIS decorator to test against torch-sparse (#9066)

pyg-team · Mar 18, 2024 · a1a3724 · a1a3724
1 parent 18d38c9
commit a1a3724
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
@@ -218,10 +218,11 @@ def withMETIS(func: Callable) -> Callable:
 
     if with_metis:
         try:  # Test that METIS can succesfully execute:
-            import pyg_lib
+            # TODO Using `pyg-lib` metis partitioning leads to some weird bugs
+            # in the # CI. As such, we require `torch-sparse` for now.
             rowptr = torch.tensor([0, 2, 4, 6])
             col = torch.tensor([1, 2, 0, 2, 1, 0])
-            pyg_lib.partition.metis(rowptr, col, num_partitions=2)
+            torch.ops.torch_sparse.partition(rowptr, col, None, 2, True)
         except Exception:
             with_metis = False