pytorch
diff --git a/‎benchmarks/dynamo/common.py
Lines changed: 5 additions & 3 deletions b/‎benchmarks/dynamo/common.py
Lines changed: 5 additions & 3 deletions
diff --git a/‎benchmarks/dynamo/torchbench.py
Lines changed: 0 additions & 22 deletions b/‎benchmarks/dynamo/torchbench.py
Lines changed: 0 additions & 22 deletions
@@ -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
 
@@ -26,20 +26,6 @@
 torch.backends.cuda.matmul.allow_tf32 = True
 
 
-def _reassign_parameters(model):
-    # torch_geometric models register parameter as tensors due to
-    # https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/dense/linear.py#L158-L168
-    # Since it is unusual thing to do, we just reassign them to parameters
-    def state_dict_hook(module, destination, prefix, local_metadata):
-        for name, param in module.named_parameters():
-            if isinstance(destination[name], torch.Tensor) and not isinstance(
-                destination[name], torch.nn.Parameter
-            ):
-                destination[name] = torch.nn.Parameter(destination[name])
-
-    model._register_state_dict_hook(state_dict_hook)
-
-
 def setup_torchbench_cwd():
     original_dir = abspath(os.getcwd())
 
@@ -314,14 +300,6 @@ def load_model(
                 extra_args=extra_args,
             )
         model, example_inputs = benchmark.get_module()
-        if model_name in [
-            "basic_gnn_edgecnn",
-            "basic_gnn_gcn",
-            "basic_gnn_sage",
-            "basic_gnn_gin",
-        ]:
-            _reassign_parameters(model)
-
         # Models that must be in train mode while training
         if is_training and (
             not use_eval_mode or model_name in self._config["only_training"]