Fixing test_load_from_yolov5

zhiqwang · Oct 10, 2021 · 5750b3c · 5750b3c
1 parent 6c21a1c
commit 5750b3c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/test/test_models.py b/test/test_models.py
@@ -316,7 +316,8 @@ def test_load_from_yolov5(arch, version, hash_prefix):
             hash_prefix=hash_prefix,
         )
 
-    model_yolov5 = YOLOv5.load_from_yolov5(yolov5s_r40_path, version=version)
+    up_version = version.replace("v", "r")
+    model_yolov5 = YOLOv5.load_from_yolov5(yolov5s_r40_path, version=up_version)
     model_yolov5.eval()
     out_from_yolov5 = model_yolov5.predict(img_path)
     assert isinstance(out_from_yolov5[0], dict)

diff --git a/yolort/models/yolo_module.py b/yolort/models/yolo_module.py
@@ -313,7 +313,7 @@ def load_from_yolov5(
             version (str): upstream version released by the ultralytics/yolov5, Possible
                 values are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0".
         """
-        model_info = load_from_ultralytics(checkpoint_path)
+        model_info = load_from_ultralytics(checkpoint_path, version=version)
         arch = f"yolov5_darknet_pan_{model_info['size']}_{version.replace('.', '')}"
         yolov5 = cls(
             lr=lr,