diff --git a/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py b/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py index 236a9328958..1ebdb434252 100644 --- a/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py +++ b/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py @@ -41,7 +41,7 @@ class RollWithEmptyAxesReplacer(FrontReplacementPattern): # reshape to original shape shape_of = Shape(graph, {'name': node_name + '/shape_of'}).create_node() - roll_node.in_port(0).get_connection().add_destination(shape_of.in_port(0)) + reshape_to_1d.in_port(0).get_connection().add_destination(shape_of.in_port(0)) reshape_to_orig_shape = Reshape(graph, {}).create_node() rename_nodes([(roll_node, node_name + '/roll'), (reshape_to_orig_shape, node_name)]) shape_of.out_port(0).connect(reshape_to_orig_shape.in_port(1)) diff --git a/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py b/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py index e172737f15e..8e60ad3ac4b 100644 --- a/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py +++ b/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py @@ -50,3 +50,8 @@ class RollWithEmptyAxesReplacerTest(unittest.TestCase): (flag, resp) = compare_graphs(graph, graph_ref, 'result', check_op_attrs=True) self.assertTrue(flag, resp) + + shape_of_nodes = graph.get_op_nodes(type='ShapeOf') + self.assertTrue(len(shape_of_nodes) == 1) + shape_of = shape_of_nodes[0] + self.assertTrue(shape_of.in_node().soft_get('name') == 'placeholder')