Update remove_converts pass with shape inference (openvinotoolkit#10474)

yli147 · Feb 17, 2022 · a090abb · a090abb
1 parent 6e5eb87
commit a090abb
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/tools/pot/openvino/tools/pot/graph/passes.py b/tools/pot/openvino/tools/pot/graph/passes.py
@@ -938,17 +938,18 @@ def find_shape_subgraph_endpoints(out_ports: List[Port], visited: set = None) ->
 
 
 def remove_converts(graph: Graph):
-    for op in graph.get_op_nodes(type='Convert'):
-        source_op = op.in_port(0).get_source().node
-        if source_op.type == 'Const' and source_op.data_type == np.float16:
-            # Get access to data node after Convert operation and set Insert_Convert_operation_after
-            # to restore Convert operation later
-            op.out_node(0)['Insert_Convert_operation_after'] = True
-            # Mark Const and Convert operation to fold them
-            source_op['need_shape_inference'] = True
-            op.out_node(0)['old_rt_info'] = op['rt_info']
-            op['stop_value_propagation'] = False
-            op['need_shape_inference'] = True
+    for op in graph.get_op_nodes():
+        if op.type == 'Convert':
+            source_op = op.in_port(0).get_source().node
+            if source_op.type == 'Const' and source_op.data_type == np.float16:
+                # Get access to data node after Convert operation and set Insert_Convert_operation_after
+                # to restore Convert operation later
+                op.out_node(0)['Insert_Convert_operation_after'] = True
+                # Mark Const and Convert operation to fold them
+                source_op['need_shape_inference'] = True
+                op.out_node(0)['old_rt_info'] = op['rt_info']
+                op['stop_value_propagation'] = False
+        op['need_shape_inference'] = True
     graph.clean_up()