[GraphBolt] Fix deprecation warning (#7493)

dmlc · Jul 1, 2024 · f9a34b6 · f9a34b6
1 parent 65b949f
commit f9a34b6
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/python/dgl/graphbolt/base.py b/python/dgl/graphbolt/base.py
@@ -79,9 +79,15 @@ def isin(elements, test_elements):
 
 if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
 
-    @torch.library.impl_abstract("graphbolt::expand_indptr")
-    def expand_indptr_abstract(indptr, dtype, node_ids, output_size):
-        """Abstract implementation of expand_indptr for torch.compile() support."""
+    torch_fake_decorator = (
+        torch.library.impl_abstract
+        if TorchVersion(torch.__version__) < TorchVersion("2.3.1")
+        else torch.library.register_fake
+    )
+
+    @torch_fake_decorator("graphbolt::expand_indptr")
+    def expand_indptr_fake(indptr, dtype, node_ids, output_size):
+        """Fake implementation of expand_indptr for torch.compile() support."""
         if output_size is None:
             output_size = torch.library.get_ctx().new_dynamic_size()
         if dtype is None: