diff --git a/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py b/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py index c0c759bd9d3..5d44e706afc 100644 --- a/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py +++ b/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py @@ -154,7 +154,6 @@ def transform_fn(data_item): calibration_dataset = nncf.Dataset(val_dataset, transform_fn) tf_quantized_model = nncf.quantize(tf_model, calibration_dataset) -tf_quantized_model = nncf.strip(tf_quantized_model) tf_quantized_model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5), @@ -168,7 +167,8 @@ def transform_fn(data_item): # Benchmark performance, calculate compression rate and validate accuracy ov_model = ov.convert_model(tf_model, share_weights=False) -ov_quantized_model = ov.convert_model(tf_quantized_model, share_weights=False) +stripped_model = nncf.strip(tf_quantized_model) +ov_quantized_model = ov.convert_model(stripped_model, share_weights=False) fp32_ir_path = ROOT / "mobilenet_v2_fp32.xml" ov.save_model(ov_model, fp32_ir_path, compress_to_fp16=False) diff --git a/nncf/common/strip.py b/nncf/common/strip.py index 8fe2cc930cd..eac68dd5f93 100644 --- a/nncf/common/strip.py +++ b/nncf/common/strip.py @@ -41,7 +41,7 @@ def strip(model: TModel, do_copy: bool = True) -> TModel: return strip_pt(model, do_copy) # type: ignore elif model_backend == BackendType.TENSORFLOW: - from nncf.tensorflow import strip as strip_tf + from nncf.tensorflow.strip import strip as strip_tf return strip_tf(model, do_copy)