fix

openvinotoolkit · Jan 14, 2025 · f993534 · f993534
1 parent c1faf40
commit f993534
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git 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
@@ -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)