From a9dc0c2c29c45e1b4a7bb764c9d6f6230750b01a Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 30 Jun 2020 15:25:38 -0700
Subject: [PATCH] update check_img_size() for model strides

---
 detect.py | 2 +-
 test.py   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/detect.py b/detect.py
index 988a002..bb84a0d 100644
--- a/detect.py
+++ b/detect.py
@@ -25,6 +25,7 @@ def detect(save_img=False):
     # torch.save(torch.load(weights, map_location=device), weights)  # update model if SourceChangeWarning
     # model.fuse()
     model.to(device).eval()
+    imgsz = check_img_size(imgsz, s=model.model[-1].stride.max())  # check img_size
     if half:
         model.half()  # to FP16
 
@@ -152,7 +153,6 @@ if __name__ == '__main__':
     parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
     parser.add_argument('--augment', action='store_true', help='augmented inference')
     opt = parser.parse_args()
-    opt.img_size = check_img_size(opt.img_size)
     print(opt)
 
     with torch.no_grad():
diff --git a/test.py b/test.py
index c0bda5f..259d444 100644
--- a/test.py
+++ b/test.py
@@ -34,6 +34,7 @@ def test(data,
         torch_utils.model_info(model)
         model.fuse()
         model.to(device)
+        imgsz = check_img_size(imgsz, s=model.model[-1].stride.max())  # check img_size
 
         # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
         # if device.type != 'cpu' and torch.cuda.device_count() > 1:
@@ -242,7 +243,6 @@ if __name__ == '__main__':
     parser.add_argument('--merge', action='store_true', help='use Merge NMS')
     parser.add_argument('--verbose', action='store_true', help='report mAP by class')
     opt = parser.parse_args()
-    opt.img_size = check_img_size(opt.img_size)
     opt.save_json = opt.save_json or opt.data.endswith('coco.yaml')
     opt.data = check_file(opt.data)  # check file
     print(opt)