Skip to content

Commit

Permalink
Fix: missed input_shape to convert_model for PT backend. (#2178)
Browse files Browse the repository at this point in the history
### Changes

Add missed `input_shape` to `convert_model` for PT backend.

### Reason for changes

IR of quantized models have dynamic shape.
  • Loading branch information
AlexanderDokuchaev authored Oct 5, 2023
1 parent 849d228 commit 2d5ea2c
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion tests/post_training/pipelines/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ def __init__(
self.model_hf = None
self.calibration_dataset = None
self.dummy_tensor = None
self.input_size = None

self.run_info = RunInfo(model=reported_name, backend=self.backend)

Expand Down Expand Up @@ -212,7 +213,7 @@ def save_quantized_model(self) -> None:
if self.backend == BackendType.OPTIMUM:
self.path_quantized_ir = self.output_model_dir / "openvino_model.xml"
elif self.backend in PT_BACKENDS:
ov_model = convert_model(self.quantized_model, example_input=self.dummy_tensor)
ov_model = convert_model(self.quantized_model, example_input=self.dummy_tensor, input_shape=self.input_size)
self.path_quantized_ir = self.output_model_dir / "model.xml"
ov.serialize(ov_model, self.path_quantized_ir)
elif self.backend == BackendType.ONNX:
Expand Down

0 comments on commit 2d5ea2c

Please sign in to comment.