fix slicing before topk

apache · Jun 2, 2021 · 0bba350 · 0bba350
1 parent e66923f
commit 0bba350
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
@@ -842,10 +842,9 @@ def true_branch():
         def false_branch():
             if isinstance(max_output_boxes_per_class, int):
                 # Do topk on smaller input if possible
-                # TODO(masahi): use axes argument in strided slice when it becomes available
-                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                slice_mx = _op.const([max_output_boxes_per_class * num_class], dtype="int64")
                 selected_scores_slice = _op.strided_slice(
-                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx
+                    selected_scores, begin=_op.const([0], dtype="int64"), end=slice_mx, axes=[1]
                 )
             else:
                 selected_scores_slice = selected_scores