Revert "Make torch_geometric models compatible with export (pytorch#123403)" (pytorch#128377)

chunyuan-w · TharinduRusira · commit dba29d2a1954 · 2024-06-14T03:54:50.000Z
This reverts commit d78991a. This PR reverts pytorch#123403 to fix the performance regression as discussed in pytorch#127513 (comment). Pull Request resolved: pytorch#128377 Approved by: https://github.com/jgong5, https://github.com/angelayi, https://github.com/desertfire
diff --git a/benchmarks/dynamo/common.py b/benchmarks/dynamo/common.py
@@ -1184,12 +1184,14 @@ def load(cls, model, example_inputs, device):
             else:
                 _register_dataclass_output_as_pytree(example_outputs)
 
-            gm = torch.export._trace._export(
+            # TODO(angelayi): change this to predispatch
+            # https://github.com/pytorch/pytorch/issues/127513 needs to be fixed before changing
+            # to predispatch to avoid performance regressions
+            gm = torch.export._trace._export_to_torch_ir(
                 model,
                 example_args,
                 example_kwargs,
-                pre_dispatch=True,
-            ).module()
+            )
             with torch.no_grad():
                 so_path = torch._inductor.aot_compile(
                     gm, example_args, example_kwargs
