minor

Author: andrey-churkin

tests/post_training/pipelines/image_classification_timm.py

@@ -45,8 +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,

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