From 39b4342c84ae5620aa244ed307b960fc1e23ea26 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Sun, 6 Mar 2022 12:16:20 +0800 Subject: [PATCH] Minor fix for weights_name --- yolort/models/yolo_lite.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolort/models/yolo_lite.py b/yolort/models/yolo_lite.py index 4ad5bb06..627f0863 100644 --- 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