Fix consts (after code-review #1)

Signed-off-by: Roman Kazantsev <[email protected]>
rkazants · Aug 27, 2020 · 41ceaa1 · 41ceaa1
1 parent 36a16b3
commit 41ceaa1
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/model-optimizer/extensions/front/tf/CTCGreedyDecoderReplacement.py b/model-optimizer/extensions/front/tf/CTCGreedyDecoderReplacement.py
@@ -52,13 +52,13 @@ def pattern(**kwargs):
                 ('shape', dict(op='ShapeOf')),
                 ('shape_1', dict(op='ShapeOf')),
                 ('strided_slice', dict(op='StridedSlice')),
-                ('stack', dict(op='Const', value=[1])),
-                ('stack1', dict(op='Const', value=[2])),
-                ('stack2', dict(op='Const', value=[1])),
+                ('stack', dict(op='Const', value=lambda v: v is not None and np.array_equal(v, [1]))),
+                ('stack1', dict(op='Const', value=lambda v: v is not None and np.array_equal(v, [2]))),
+                ('stack2', dict(op='Const', value=lambda v: v is not None and np.array_equal(v, [1]))),
                 ('strided_slice_1', dict(op='StridedSlice')),
-                ('stack_1', dict(op='Const', value=[0])),
-                ('stack1_1', dict(op='Const', value=[1])),
-                ('stack2_1', dict(op='Const', value=[1])),
+                ('stack_1', dict(op='Const', value=lambda v: v is not None and np.array_equal(v, [0]))),
+                ('stack1_1', dict(op='Const', value=lambda v: v is not None and np.array_equal(v, [1]))),
+                ('stack2_1', dict(op='Const', value=lambda v: v is not None and np.array_equal(v, [1]))),
                 ('dims', dict(op='Pack')),
                 ('fill', dict(op='Fill')),
                 ('decoder', dict(op='CTCGreedyDecoder')),