fix batch inference bug (#413)

2022-12-29 17:18:54 +08:00 · 2022-12-29 17:18:54 +08:00 · cdc359c2de
parent 99af3c980f
commit cdc359c2de
1 changed files with 1 additions and 1 deletions
--- a/mmyolo/models/dense_heads/yolov5_head.py
+++ b/mmyolo/models/dense_heads/yolov5_head.py
@ -365,7 +365,7 @@ class YOLOv5Head(BaseDenseHead):
            ]
            flatten_objectness = torch.cat(flatten_objectness, dim=1).sigmoid()
        else:
-            flatten_objectness = [None for _ in range(len(featmap_sizes))]
+            flatten_objectness = [None for _ in range(num_imgs)]

        results_list = []
        for (bboxes, scores, objectness,