Merge pull request #829 from pfnet/graph-context-interface-update

update new GraphContext interface
pfnet · Oct 1, 2024 · 0e842f7 · 0e842f7
2 parents c493274 + 68f8fba
commit 0e842f7
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/pytorch_pfn_extras/onnx/pfto_exporter/export.py b/pytorch_pfn_extras/onnx/pfto_exporter/export.py
@@ -1,4 +1,5 @@
 import dataclasses
+from functools import partial
 import types
 import typing
 import warnings
@@ -115,7 +116,11 @@
     "special_gammaln": "lgamma",
 }
 
-if pytorch_pfn_extras.requires("1.13"):
+GraphContext: Any
+if pytorch_pfn_extras.requires("2.4"):
+    from torch.onnx._internal import jit_utils
+    GraphContext = partial(jit_utils.GraphContext, values_in_env=set())
+elif pytorch_pfn_extras.requires("1.13"):
     from torch.onnx._internal import jit_utils
     GraphContext = jit_utils.GraphContext
 else: