Minor fix for weights_name

zhiqwang · Mar 6, 2022 · 39b4342 · 39b4342
1 parent 6ee9f36
commit 39b4342
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/yolort/models/yolo_lite.py b/yolort/models/yolo_lite.py
@@ -143,7 +143,7 @@ def _yolov5_mobilenet_v3_small_fpn(
     if pretrained:
         if model_urls.get(weights_name, None) is None:
             raise ValueError(f"No checkpoint is available for model {weights_name}")
-        state_dict = load_state_dict_from_url(model_urls["retinanet_resnet50_fpn_coco"], progress=progress)
+        state_dict = load_state_dict_from_url(model_urls[weights_name], progress=progress)
         model.load_state_dict(state_dict)
     return model