From b70e26c73e277504a8f24799f2bad107f8ce8ffe Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 21 Dec 2020 07:30:06 +0900 Subject: [PATCH] fix topi vision test by wrapping tir const around int argument --- python/tvm/topi/cuda/nms.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/tvm/topi/cuda/nms.py b/python/tvm/topi/cuda/nms.py index c71ea4d728a20..8a372f2cabd09 100644 --- a/python/tvm/topi/cuda/nms.py +++ b/python/tvm/topi/cuda/nms.py @@ -547,6 +547,9 @@ def nms_inner_loop(ib, j): box_indices[i, num_valid_boxes_local[0]] = indices[i, orig_idx] num_valid_boxes_local[0] += 1 + if isinstance(max_output_size, int): + max_output_size = tvm.tir.const(max_output_size) + with ib.if_scope(tvm.tir.all(iou_threshold > 0, valid_count[i] > 0)): # Apply nms with ib.for_range(0, valid_count[i]) as j: