merge fix to export test

sony · Jan 2, 2025 · 34985e9 · 34985e9
1 parent 1d6f511
commit 34985e9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/...compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/latest/__init__.py b/...compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/latest/__init__.py
@@ -18,8 +18,8 @@
 if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model as get_keras_tpc_latest
     from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
-        get_tpc_model as generate_keras_tpc, get_tpc_model as generate_pytorch_tpc
+        get_tpc_model as generate_keras_tpc
 if FOUND_TORCH:
     from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model as get_pytorch_tpc_latest
     from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
-        get_tpc_model as generate_pytorch_tpc, get_tpc_model as generate_pytorch_tpc
+        get_tpc_model as generate_pytorch_tpc
diff --git a/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tp_model.py b/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tp_model.py
@@ -180,6 +180,7 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_UNSQUEEZE.value, qc_options=no_quantization_config))
     operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_SIZE.value, qc_options=no_quantization_config))
     operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_PERMUTE.value, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_BATCH_NORM.value, qc_options=no_quantization_config))
     operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_TRANSPOSE.value, qc_options=no_quantization_config))
     operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_EQUAL.value, qc_options=no_quantization_config))
     operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.OPSET_ARGMAX.value, qc_options=no_quantization_config))

diff --git a/tests/pytorch_tests/exporter_tests/test_exporting_qat_models.py b/tests/pytorch_tests/exporter_tests/test_exporting_qat_models.py
@@ -21,7 +21,6 @@
 from torchvision.models.mobilenetv2 import mobilenet_v2
 
 import model_compression_toolkit as mct
-from model_compression_toolkit.verify_packages import FOUND_ONNXRUNTIME, FOUND_ONNX
 from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \