minor

Author: andrey-churkin

tests/post_training/pipelines/image_classification_timm.py

@@ -45,9 +45,8 @@ def prepare_model(self) -> None:
             onnx_path = self.fp32_model_dir / "model_fp32.onnx"
             additional_kwargs = {}
             if self.batch_size > 1:
-                batch = torch.export.Dim("batch")
-                additional_kwargs["dynamic_shapes"] = ({0: batch},)
-
+                additional_kwargs["input_names"] = ["image"]
+                additional_kwargs["dynamic_axes"] = {"image": {0: "batch"}}
             torch.onnx.export(
                 timm_model,
                 self.dummy_tensor,
@@ -56,7 +55,6 @@ def prepare_model(self) -> None:
                 opset_version=13,
                 **additional_kwargs,
             )
-
             self.model = onnx.load(onnx_path)
             self.input_name = self.model.graph.input[0].name
 

tests/post_training/requirements.txt

@@ -22,4 +22,4 @@ accelerate==1.9.0
 transformers==4.53.0
 whowhatbench @ git+https://github.com/openvinotoolkit/[email protected]#subdirectory=tools/who_what_benchmark
 datasets==3.6.0
-onnxscript
+onnxscript==0.5.4