minimum fixed

neo-ai · May 29, 2021 · d203562 · d203562
1 parent 3fe91e8
commit d203562
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 5 deletions.
diff --git a/python/tvm/topi/cuda/nms.py b/python/tvm/topi/cuda/nms.py
@@ -1166,8 +1166,6 @@ def all_class_non_max_suppression(
     topk_indices = expand_dims(topk_indices, axis=0)
     final_indices = gather_nd(selected_indices, topk_indices, batch_dims=1)
 
-    # TODO
-    # num_detections = minimum(num_total_detections, max_detection_per_batch)
-    num_detections = num_total_detections
+    num_detections = minimum(num_total_detections, max_detection_per_batch)
 
     return [final_indices, num_detections]
diff --git a/python/tvm/topi/cuda/vision.py b/python/tvm/topi/cuda/vision.py
@@ -39,7 +39,9 @@ def traverse(op):
                 traverse(tensor.op)
         scheduled_ops.append(op)
 
-    traverse(outs[0].op)
+    for o in outs:
+        traverse(o.op)
+
     return s
 
 

diff --git a/python/tvm/topi/vision/nms_util.py b/python/tvm/topi/vision/nms_util.py
@@ -180,7 +180,6 @@ def collect_selected_indices_tf(num_class, selected_indices, num_detections, row
     )
 
 
-
 def _all_class_nms_ir(
     boxes,
     sorted_scores,