From ef00498385d288dea544a4a95d2b98b25a6fea3c Mon Sep 17 00:00:00 2001 From: Krzysztof Czugala Date: Mon, 8 Aug 2022 08:55:00 +0200 Subject: [PATCH] [MO] Result edges naming issue (#12154) * the first approach to fixing the Result edge naming issue * more robust approach, not depending on the split * out_nodes_ids type change * missing if clause --- tools/mo/openvino/tools/mo/front/output_cut.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tools/mo/openvino/tools/mo/front/output_cut.py b/tools/mo/openvino/tools/mo/front/output_cut.py index cf7207ddeb9..2d091f8160f 100644 --- a/tools/mo/openvino/tools/mo/front/output_cut.py +++ b/tools/mo/openvino/tools/mo/front/output_cut.py @@ -27,6 +27,7 @@ class OutputCut(FrontReplacementPattern): for node in graph.get_op_nodes(needs_removal=True): fw_info = None in_node = None + out_nodes_ids = {} for in_port_idx in node.in_edges(): node_idx = node.in_edge(in_port_idx)['in'] if node_idx in node.in_nodes(): @@ -35,9 +36,15 @@ class OutputCut(FrontReplacementPattern): if fw_info_value: fw_info = fw_info_value break + if fw_info is not None and in_node is not None: + for out_idx in in_node.out_nodes(): + out_node = in_node.out_node(out_idx) + out_nodes_ids[out_idx] = out_node.id + graph.erase_node(node) if fw_info is not None and in_node is not None: for out_idx in in_node.out_nodes(): - set_edge_attribute_between_nodes(in_node, in_node.out_node(out_idx), + if node.id == out_nodes_ids[out_idx]: + set_edge_attribute_between_nodes(in_node, in_node.out_node(out_idx), 'fw_tensor_debug_info', fw_info)