From f53124e7381d0d86c38d96c57393de796a66c4a9 Mon Sep 17 00:00:00 2001 From: Evgenya Stepyreva Date: Tue, 25 Aug 2020 20:30:41 +0300 Subject: [PATCH] [ MO ] Fix for remove_op_node_with_data_node (#1934) --- model-optimizer/extensions/middle/UselessMerge.py | 2 +- model-optimizer/mo/middle/passes/eliminate.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/model-optimizer/extensions/middle/UselessMerge.py b/model-optimizer/extensions/middle/UselessMerge.py index e03cd5f6936f7e..7846c49cf38818 100644 --- 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 index f0b373b5b2dc1b..ba396243abf324 100644 --- 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])