Fix tuning failure after stripping equivalent nodes (intel#1290)

chensuyue · Sep 27, 2022 · 0791433 · 0791433
1 parent 3827826
commit 0791433
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/generic/insert_print_node.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/generic/insert_print_node.py
@@ -151,6 +151,8 @@ def do_transformation(self):
                 if post_node_names:
                     for post_node_name in post_node_names:
                         post_node = graph_info[post_node_name].node
+                        if each_node_name not in post_node.input:
+                            continue
                         if post_node.op == 'FusedBatchNormV3':
                             if "_print_identity" in \
                                graph_info[Helper.node_name_from_input(post_node.name)].node.input[0]:

diff --git a/neural_compressor/adaptor/tf_utils/util.py b/neural_compressor/adaptor/tf_utils/util.py
@@ -401,7 +401,7 @@ def is_equivalent_input(input_tensor_list_1, input_tensor_list_2):
     for node_to_remove in nodes_to_remove:
         stripped_graph.remove_node(node_to_remove)
     return tf.compat.v1.graph_util.extract_sub_graph \
-        (stripped_graph.dump_graph(), set(stripped_graph_node_names).intersection(output_node_names)), \
+        (stripped_graph.dump_graph(), list(set(stripped_graph_node_names).intersection(output_node_names))), \
             replaced_nodes_type
 
 # THIS API IS TO BE DEPRECATED!