From 954bc6fedf128fce6dae80e20dd22ca86eabd9d1 Mon Sep 17 00:00:00 2001 From: Roman Kazantsev Date: Mon, 8 Jun 2020 17:43:32 +0500 Subject: [PATCH] Fix clean-up for transformation #2 Signed-off-by: Roman Kazantsev --- model-optimizer/extensions/middle/sparse_reshape.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model-optimizer/extensions/middle/sparse_reshape.py b/model-optimizer/extensions/middle/sparse_reshape.py index 72e4174aa53e58..5db8b162dcd261 100644 --- a/model-optimizer/extensions/middle/sparse_reshape.py +++ b/model-optimizer/extensions/middle/sparse_reshape.py @@ -60,7 +60,7 @@ def replace_pattern(self, graph: Graph, match: dict): raise Error("Input shape and output shape values must be equal for node {}".format(sparse_reshape.id)) nodes_to_remove = [sparse_reshape.id] - if 0 in sparse_reshape.out_nodes(): + if sparse_reshape.is_out_port_connected(0): sparse_reshape.out_port(0).get_connection().set_source(sparse_reshape.in_port(0).get_source()) output_data_node = sparse_reshape.out_node(0) nodes_to_remove.append(output_data_node.id) @@ -68,7 +68,7 @@ def replace_pattern(self, graph: Graph, match: dict): input_data_node = sparse_reshape.in_node(0) nodes_to_remove.append(input_data_node.id) - if 1 in sparse_reshape.out_nodes(): + if sparse_reshape.is_out_port_connected(1): sparse_reshape.out_port(1).get_connection().set_source(sparse_reshape.in_port(1).get_source()) output_data_node = sparse_reshape.out_node(1) nodes_to_remove.append(output_data_node.id)