simplify

trevor-m · May 29, 2021 · 44d88cd · 44d88cd
1 parent 74e1917
commit 44d88cd
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
@@ -858,8 +858,7 @@ def false_branch():
             max_output_boxes_per_batch, int
         ), "dynamic number of boxes not supported yet."
         nmsed_scores, topk_indices = select_topk(max_output_boxes_per_batch < max_total_size)
-        topk_indices = _op.expand_dims(topk_indices, axis=0)
-        indices = _op.gather_nd(selected_indices, topk_indices, batch_dims=1)
+        indices = _op.take(selected_indices, topk_indices, axis=1, batch_dims=1)
 
         nmsed_box_indices = _op.take(indices, _op.const(1), axis=2)
         nmsed_classes = _op.cast(_op.take(indices, _op.const(0), axis=2), "float32")