From 046478e9350b96c761a4b4f1fc852888381ae2c9 Mon Sep 17 00:00:00 2001 From: BowenBao Date: Thu, 7 Dec 2023 18:39:16 -0800 Subject: [PATCH] Fix model name retrieval and remove unused argument (#115108) Summary: 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. X-link: https://github.com/pytorch/pytorch/pull/115108 Approved by: https://github.com/thiagocrepaldi Reviewed By: atalman Differential Revision: D51939505 fbshipit-source-id: ea6a78bd2fd85689ee0ba8e17c2be58da1b6faee --- userbenchmark/dynamo/dynamobench/common.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/userbenchmark/dynamo/dynamobench/common.py b/userbenchmark/dynamo/dynamobench/common.py index 96d356b8b3..28ac042a96 100644 --- a/userbenchmark/dynamo/dynamobench/common.py +++ b/userbenchmark/dynamo/dynamobench/common.py @@ -704,7 +704,6 @@ def speedup_experiment_ds(args, model_iter_fn, model, example_inputs): def speedup_experiment_onnx( - onnx_model_cls: Type[OnnxModelFromTorchScript], args, model_iter_fn, onnx_model: OnnxModel, @@ -1069,10 +1068,7 @@ class OnnxModel(abc.ABC): _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 )