Fix model_extraction_transformation

openvinotoolkit · Jul 18, 2024 · 387e666 · 387e666
1 parent d113c2a
commit 387e666
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/nncf/openvino/graph/model_transformer.py b/nncf/openvino/graph/model_transformer.py
@@ -581,10 +581,11 @@ def _apply_model_extraction_transformation(
             output_node = name_to_node_mapping[output_name]
 
             result_name = get_result_node_name(output_name, output_port_id)
-            if output_node.get_element_type() != outputs_type:
-                output_node = opset.convert(output_node, destination_type=outputs_type)
-            new_result = opset.result(output_node, name=result_name)
-            result_tensor_names = [result_name] + list(output_node.output(0).get_names())
+            output_port = output_node.output(output_port_id)
+            if output_port.get_element_type() != outputs_type:
+                output_port = opset.convert(output_node, destination_type=outputs_type).output(0)
+            new_result = opset.result(output_port, name=result_name)
+            result_tensor_names = [result_name] + list(output_port.get_names())
             OVModelTransformer._update_tensor_names([new_result.get_output_tensor(0)], result_tensor_names)
             results.append(new_result)