Support dynamic batch dimension in model converter

model_converter/src/model_converter/converters/pytorch.py

@@ -154,6 +154,7 @@ def export_to_openvino(
         input_names: list[str] | None = None,
         output_names: list[str] | None = None,
         metadata: dict[tuple[str, str], str] | None = None,
+        batch_size: int = -1,
     ) -> tuple[Path, Path]:
         """Export PyTorch model to OpenVINO format.
 
@@ -165,6 +166,9 @@ def export_to_openvino(
             input_names: Names for input tensors
             output_names: Names for output tensors
             metadata: Metadata to embed in the model
+            batch_size: Batch dimension for the exported model. Use ``-1`` (default) to
+                keep the batch dimension dynamic so any batch size works at inference,
+                or a positive integer to fix the batch dimension to that value.
 
         Returns:
             Tuple of (fp16_model_path, fp32_model_path) - FP16 for final use, FP32 for quantization
@@ -175,16 +179,20 @@ def export_to_openvino(
             model = self._prepare_model_for_export(model, model_config)
             model.eval()
             dummy_input = self._create_example_input(input_shape, model_config)
+            # Keep spatial dims static; the batch dimension follows ``batch_size``
+            # (``-1`` keeps it dynamic so any batch size works at inference).
+            target_shape = ov.PartialShape([batch_size, *input_shape[1:]])
+
             self.logger.info("Direct PyTorch to OpenVINO conversion")
-            ov_model = ov.convert_model(model, example_input=dummy_input)
+            ov_model = ov.convert_model(model, example_input=dummy_input, input=(target_shape,))
             self.logger.info("✓ PyTorch to OpenVINO conversion complete")
 
-            # Reshape model to fixed input shape (remove dynamic dimensions)
+            # Reshape model to the requested input shape (dynamic batch when batch_size == -1)
             first_input = ov_model.input(0)
             input_name_for_reshape = next(iter(first_input.get_names())) if first_input.get_names() else 0
 
-            self.logger.debug(f"Setting fixed input shape: {input_shape}")
-            ov_model.reshape({input_name_for_reshape: input_shape})
+            self.logger.debug(f"Setting input shape: {target_shape}")
+            ov_model.reshape({input_name_for_reshape: target_shape})
 
             # Post-process the model
             ov_model = self._postprocess_openvino_model(

model_converter/tests/unit/test_registry.py

@@ -447,9 +447,11 @@ def test_export_to_openvino_converts_and_saves_models(
     ):
         """export_to_openvino converts the model and saves FP32 and FP16 artifacts."""
         dummy_input = object()
+        dynamic_shape = object()
         openvino_module = types.ModuleType("openvino")
         openvino_module.convert_model = MagicMock(return_value=mock_ov_model)
         openvino_module.save_model = MagicMock()
+        openvino_module.PartialShape = MagicMock(return_value=dynamic_shape)
 
         with (
             patch.dict(sys.modules, {"openvino": openvino_module}),
@@ -473,8 +475,13 @@ def test_export_to_openvino_converts_and_saves_models(
         assert fp32_path == converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml"
         mock_prepare.assert_called_once_with(mock_torch_model, sample_model_config)
         mock_example_input.assert_called_once_with([1, 3, 224, 224], sample_model_config)
-        openvino_module.convert_model.assert_called_once_with(mock_torch_model, example_input=dummy_input)
-        mock_ov_model.reshape.assert_called_once_with({"input": [1, 3, 224, 224]})
+        openvino_module.PartialShape.assert_called_once_with([-1, 3, 224, 224])
+        openvino_module.convert_model.assert_called_once_with(
+            mock_torch_model,
+            example_input=dummy_input,
+            input=(dynamic_shape,),
+        )
+        mock_ov_model.reshape.assert_called_once_with({"input": dynamic_shape})
         mock_postprocess.assert_called_once_with(
             mock_ov_model,
             input_names=["input"],