fix yolo classify model loading error (ultralytics#9196)

hmurari · Apr 17, 2024 · 4fb5e94 · 4fb5e94
1 parent 1da7da9
commit 4fb5e94
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ultralytics/models/yolo/classify/train.py b/ultralytics/models/yolo/classify/train.py
@@ -73,7 +73,7 @@ def setup_model(self):
         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