[ MO ] Fix for remove_op_node_with_data_node (openvinotoolkit#1934)

mryzhov · Aug 26, 2020 · 3304f41 · 3304f41
1 parent 43bfb0f
commit 3304f41
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/model-optimizer/extensions/middle/UselessMerge.py b/model-optimizer/extensions/middle/UselessMerge.py
@@ -40,5 +40,5 @@ def pattern(self):
 
     def replace_pattern(self, graph: Graph, match: dict):
         if len(graph.in_edges(match['merge'].id)) <= 1:
-            remove_op_node_with_data_node(graph, match['merge'])
+            remove_op_node_with_data_node(graph, match['merge'], list(match['merge'].in_nodes().values())[0])
             log.info("Useles Merge op and data nodes was deleted op='{}'".format(match['merge'].id))
diff --git a/model-optimizer/mo/middle/passes/eliminate.py b/model-optimizer/mo/middle/passes/eliminate.py
@@ -217,10 +217,11 @@ def merge_data_nodes(graph, survived, removed):
 
 
 # TODO: unit tests
-def remove_op_node_with_data_node(graph, node_to_remove):
+def remove_op_node_with_data_node(graph, node_to_remove, input_data_node=None):
     from mo.graph.graph import Node
     assert node_to_remove.kind == 'op'
-    input_data_node = node_to_remove.in_node()
+    if input_data_node is None:
+        input_data_node = node_to_remove.in_node()
     output_node = [v for _, v in graph.out_edges(node_to_remove.id)]
     assert len(output_node) == 1, "Cannot remove node producing two or more output tensors"
     output_node = Node(graph, output_node[0])