fix topi vision test by wrapping tir const around int argument

apache · Dec 20, 2020 · b70e26c · b70e26c
1 parent 5179cf8
commit b70e26c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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: