Fix after review #4

Signed-off-by: Roman Kazantsev <[email protected]>
rkazants · Aug 28, 2020 · 87026e3 · 87026e3
1 parent 5027645
commit 87026e3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/model-optimizer/extensions/front/tf/CTCGreedyDecoderReplacement_test.py b/model-optimizer/extensions/front/tf/CTCGreedyDecoderReplacement_test.py
@@ -63,7 +63,8 @@ def test1(self):
                                 nodes_with_edges_only=True)
 
         (flag, resp) = compare_graphs(graph, graph_ref, 'last', check_op_attrs=True)
-        self.assertEqual(graph.get_op_nodes(op='Cast')[0]['name'], 'sparse_to_dense',
+        self.assertEqual(len(graph.get_op_nodes(op='Cast')) == 1 and
+                         graph.get_op_nodes(op='Cast')[0]['name'] == 'sparse_to_dense', True,
                          'Name is not inherited from original node for CTCGreedyDecoderReplacement')
         self.assertTrue(flag, resp)
 
@@ -158,6 +159,7 @@ def test2(self):
                                 nodes_with_edges_only=True)
 
         (flag, resp) = compare_graphs(graph, graph_ref, 'last', check_op_attrs=True)
-        self.assertEqual(graph.get_op_nodes(op='Cast')[0]['name'], 'sparse_to_dense',
+        self.assertEqual(len(graph.get_op_nodes(op='Cast')) == 1 and
+                         graph.get_op_nodes(op='Cast')[0]['name'] == 'sparse_to_dense', True,
                          'Name is not inherited from original node for CTCGreedyDecoderReplacement2')
         self.assertTrue(flag, resp)