initialization bug fixed in cuda

trevor-m · May 29, 2021 · c5718e2 · c5718e2
1 parent 5623e3f
commit c5718e2
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/python/tvm/topi/cuda/nms.py b/python/tvm/topi/cuda/nms.py
@@ -1022,6 +1022,18 @@ def _collect_selected_indices_and_scores_ir(
     ib.scope_attr(by, "thread_extent", nthread_by)
     zero = cast(0, "int64")
 
+    with ib.new_scope():
+        idx = bx * nthread_tx + tx
+        idy = cast(by, "int64")
+        batch_id = idy // num_class
+        class_id = idy % num_class
+
+        with ib.if_scope(idx < num_boxes):
+            offset = idx + class_id * num_boxes
+            collected_indices[batch_id, offset, 0] = zero
+            collected_indices[batch_id, offset, 1] = zero
+            collected_scores[batch_id, offset] = -1.0
+
     with ib.new_scope():
         idx = bx * nthread_tx + tx
         idy = cast(by, "int64")
@@ -1033,11 +1045,6 @@ def _collect_selected_indices_and_scores_ir(
             collected_indices[batch_id, offset, 0] = class_id
             collected_indices[batch_id, offset, 1] = cast(selected_indices[idy, idx], "int64")
             collected_scores[batch_id, offset] = selected_scores[idy, idx]
-        with ib.else_scope():
-            with ib.if_scope(idx < num_boxes):
-                collected_indices[batch_id, offset, 0] = zero
-                collected_indices[batch_id, offset, 1] = zero
-                collected_scores[batch_id, offset] = -1.0
 
     return ib.get()