Skip to content

Commit

Permalink
Add NAS autodownload (ultralytics#14627)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored Jul 23, 2024
1 parent f94c82d commit 82956dc
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion ultralytics/models/nas/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import torch

from ultralytics.engine.model import Model
from ultralytics.utils.downloads import attempt_download_asset
from ultralytics.utils.torch_utils import model_info, smart_inference_mode

from .predict import NASPredictor
Expand Down Expand Up @@ -56,7 +57,7 @@ def _load(self, weights: str, task: str):

suffix = Path(weights).suffix
if suffix == ".pt":
self.model = torch.load(weights)
self.model = torch.load(attempt_download_asset(weights))
elif suffix == "":
self.model = super_gradients.training.models.get(weights, pretrained_weights="coco")
# Standardize model
Expand Down

0 comments on commit 82956dc

Please sign in to comment.