From 1b3d00bc64b099ab2f59d633215ca2880888b7cd Mon Sep 17 00:00:00 2001 From: masahi Date: Fri, 12 Mar 2021 00:56:28 +0900 Subject: [PATCH] [ONNX] Use take instead of min in NMS conditions (#7633) --- python/tvm/relay/frontend/onnx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 860753d6cd0b..f31b8c927f8f 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2453,7 +2453,7 @@ def _first_cond( nms_size_out, ): # Loop over classes, end when i == C - return _op.min(_op.less(i, C)) + return _op.take(_op.less(i, C), _expr.const(0)) def _first_body( i, @@ -2561,7 +2561,7 @@ def _first_body( def _inner_cond(i, j, C, onnx_out, nms_size, out): # inner loop over number of classes - return _op.min(_op.less(j, C)) + return _op.take(_op.less(j, C), _expr.const(0)) def _inner_body(i, j, C, onnx_out, nms_size, out): # slice to get current batch and class for valid box indicator @@ -2591,7 +2591,7 @@ def _inner_body(i, j, C, onnx_out, nms_size, out): def _outer_cond(i, B, C, onnx_out, nms_size_out, out): # Outer loop is over batch size - return _op.min(_op.less(i, B)) + return _op.take(_op.less(i, B), _expr.const(0)) def _outer_body(i, B, C, onnx_out, nms_size_out, out): # Outer loop just calls inner loop