Fix torch_geometric.compile for PyTorch < 2.1 (#8240)

pyg-team · Oct 21, 2023 · 7085a3a · 7085a3a
1 parent cd6d1f2
commit 7085a3a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torch_geometric/compile.py b/torch_geometric/compile.py
@@ -93,7 +93,10 @@ def fn(model: Callable) -> Callable:
     model = to_jittable(model)
 
     # Do not generate device asserts which may slow down model execution:
-    torch._inductor.config.triton.assert_indirect_indexing = False
+    try:
+        torch._inductor.config.triton.assert_indirect_indexing = False
+    except AttributeError:  # PyTorch < 2.1:
+        pass
 
     # Finally, run `torch.compile` to create an optimized version:
     out = torch.compile(model, *args, **kwargs)