diff --git a/test.py b/test.py
index f218e7e9..7d9f8822 100644
--- a/test.py
+++ b/test.py
@@ -93,7 +93,7 @@ def test(
                 # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
                 image_id = int(Path(paths[si]).stem.split('_')[-1])
                 box = pred[:, :4].clone()  # xyxy
-                scale_coords(imgs[si].shape, box, shapes[si])  # to original shape
+                scale_coords(imgs[si].shape[1:], box, shapes[si])  # to original shape
                 box = xyxy2xywh(box)  # xywh
                 box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
                 for di, d in enumerate(pred):