Fix for comments

apache · Jul 21, 2023 · 2f8ae58 · 2f8ae58
1 parent 80aa621
commit 2f8ae58
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/python/tvm/relay/op/vision/_vision.py b/python/tvm/relay/op/vision/_vision.py
@@ -142,7 +142,7 @@ def _regular_nms_shape_func(boxes_shape, scores_shape, attrs):
     out_scores_shape[0] = boxes_shape[0]
     out_scores_shape[1] = int64(attrs.max_detections)
 
-    out_num_detections_shape[0] = int64(1)
+    out_num_detections_shape[0] = boxes_shape[0]
 
     return out_boxes_shape, out_classes_shape, out_scores_shape, out_num_detections_shape
 

diff --git a/python/tvm/topi/cuda/ssd/multibox.py b/python/tvm/topi/cuda/ssd/multibox.py
@@ -363,9 +363,7 @@ def transform_loc(loc, loc_base_idx, anchor, anchor_base_idx, clip, vx, vy, vw,
         with ib.if_scope(tvm.tir.any(keep_background == 1, cls_id[tid] > 0)):
             with ib.if_scope(j == 0):
                 out_base_idx = i * num_anchors * 6
-                out_loc[out_base_idx] = (
-                    cls_id[tid] - 0.0 if keep_background == 1 else cls_id[tid] - 1.0
-                )
+                out_loc[out_base_idx] = cls_id[tid] if keep_background == 1 else cls_id[tid] - 1.0
                 out_loc[out_base_idx + 1] = score[tid]
                 (
                     out_loc[out_base_idx + 2],
@@ -385,9 +383,7 @@ def transform_loc(loc, loc_base_idx, anchor, anchor_base_idx, clip, vx, vy, vw,
                 )
             with ib.else_scope():
                 out_base_idx = i * num_anchors * 6 + temp_valid_count[tid - 1] * 6
-                out_loc[out_base_idx] = (
-                    cls_id[tid] - 0.0 if keep_background == 1 else cls_id[tid] - 1.0
-                )
+                out_loc[out_base_idx] = cls_id[tid] if keep_background == 1 else cls_id[tid] - 1.0
                 out_loc[out_base_idx + 1] = score[tid]
                 (
                     out_loc[out_base_idx + 2],