From 5addd81b04da2a866d171f7b5d746b9bb7e7e8b3 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 11 Mar 2021 16:18:39 +0900 Subject: [PATCH] [ONNX] Use take instead of min in NMS conditions --- 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