[GraphBolt] Fix deprecation warning

dmlc · Jun 29, 2024 · 1bd41ff · 1bd41ff
1 parent 489671c
commit 1bd41ff
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/python/dgl/graphbolt/base.py b/python/dgl/graphbolt/base.py
@@ -79,7 +79,12 @@ def isin(elements, test_elements):
 
 if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
 
-    @torch.library.impl_abstract("graphbolt::expand_indptr")
+    @getattr(
+        torch.library,
+        "impl_abstract"
+        if TorchVersion(torch.__version__) < TorchVersion("2.3.1")
+        else "register_fake",
+    )("graphbolt::expand_indptr")
     def expand_indptr_abstract(indptr, dtype, node_ids, output_size):
         """Abstract implementation of expand_indptr for torch.compile() support."""
         if output_size is None: