fix failing unittests

Lightning-AI · Nov 24, 2021 · ff78b44 · ff78b44
1 parent ee16652
commit ff78b44
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchmetrics/detection/map.py b/torchmetrics/detection/map.py
@@ -374,7 +374,7 @@ def _get_coco_format(
         annotation_id = 1  # has to start with 1, otherwise COCOEval results are wrong
 
         boxes = [
-            box_convert(box, in_fmt="xyxy", out_fmt="xywh") if box.size() == torch.Size([1, 4]) else box
+            box_convert(box, in_fmt="xyxy", out_fmt="xywh") if box.ndim > 1 and box.size(1) == 4 else box
             for box in boxes
         ]
         for image_id, (image_boxes, image_labels) in enumerate(zip(boxes, labels)):