From 292e0287793b73a591a4a0fa2c2e252912a3fe06 Mon Sep 17 00:00:00 2001 From: Mo Li <82895469+DseidLi@users.noreply.github.com> Date: Sat, 23 Mar 2024 01:04:34 +0800 Subject: [PATCH] fix yolo classify model loading error (#9196) --- ultralytics/models/yolo/classify/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/models/yolo/classify/train.py b/ultralytics/models/yolo/classify/train.py index a7bbe772..42c65542 100644 --- a/ultralytics/models/yolo/classify/train.py +++ b/ultralytics/models/yolo/classify/train.py @@ -73,7 +73,7 @@ class ClassificationTrainer(BaseTrainer): elif model in torchvision.models.__dict__: self.model = torchvision.models.__dict__[model](weights="IMAGENET1K_V1" if self.args.pretrained else None) else: - FileNotFoundError(f"ERROR: model={model} not found locally or online. Please check model name.") + raise FileNotFoundError(f"ERROR: model={model} not found locally or online. Please check model name.") ClassificationModel.reshape_outputs(self.model, self.data["nc"]) return ckpt