fix : add if statement when n_queries * n_cls is less than 10000

jozhang97 · Aug 4, 2023 · 4ea705b · 4ea705b
1 parent 985fa0b
commit 4ea705b
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/models/deformable_detr.py b/models/deformable_detr.py
@@ -486,11 +486,12 @@ def forward(self, outputs, target_sizes):
             score = all_scores[b]
             lbls = all_labels[b]
 
-            pre_topk = score.topk(10000).indices
-            box = box[pre_topk]
-            score = score[pre_topk]
-            lbls = lbls[pre_topk]
-
+            if n_queries * n_cls > 10000:
+                pre_topk = score.topk(10000).indices
+                box = box[pre_topk]
+                score = score[pre_topk]
+                lbls = lbls[pre_topk]
+
             keep_inds = batched_nms(box, score, lbls, 0.7)[:100]
             results.append({
                 'scores': score[keep_inds],