From cdc359c2debfbaebcce3554e36a78b364fd56380 Mon Sep 17 00:00:00 2001 From: RangiLyu Date: Thu, 29 Dec 2022 17:18:54 +0800 Subject: [PATCH] fix batch inference bug (#413) --- mmyolo/models/dense_heads/yolov5_head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmyolo/models/dense_heads/yolov5_head.py b/mmyolo/models/dense_heads/yolov5_head.py index be76f48f..2e4d214a 100644 --- 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,