[ONNX][Bench] Fix model name retrieval and remove unused argument (#115108)
Might be some upstream updates, the previous hack starts to not pick up model names, updating to use the other more appropriate variable.
Also fix a bug with an unused argument that was supposed to be removed.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115108
Approved by: https://github.com/thiagocrepaldi
diff --git a/benchmarks/dynamo/common.py b/benchmarks/dynamo/common.py
index 96d356b..28ac042 100644
--- a/benchmarks/dynamo/common.py
+++ b/benchmarks/dynamo/common.py
@@ -704,7 +704,6 @@
def speedup_experiment_onnx(
- onnx_model_cls: Type[OnnxModelFromTorchScript],
args,
model_iter_fn,
onnx_model: OnnxModel,
@@ -1069,10 +1068,7 @@
_COMPILER_NAME: str
def __init__(self, output_directory, model, example_inputs, dynamic_shapes: bool):
- # Hack to get model name.
- from torch._functorch import aot_autograd
-
- model_name = aot_autograd.model_name
+ model_name = current_name
self.model_dir = self._generate_onnx_model_directory(
output_directory, self._COMPILER_NAME, model_name
)
