FIx downstream impact by the replacement of torch.load in ultralytics package (#13623)

* Update torch_load

* Auto-format by https://ultralytics.com/actions

---------

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Jing QiuandUltralyticsAssistant authored and GitHub committed 2025-07-04 12:07:46 +02:00
1 parent 79c4c31d16
commit 2540fd4c1c
9 files changed
+20 -12

No files matched your search

+2 -2
View File
@@ -10,7 +10,7 @@ Usage:
model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') # local repo
"""
import torch
from ultralytics.utils.patches import torch_load
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
@@ -87,7 +87,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
cfg = list((Path(__file__).parent / "models").rglob(f"{path.stem}.yaml"))[0] # model.yaml path
model = DetectionModel(cfg, channels, classes) # create model
if pretrained:
ckpt = torch.load(attempt_download(path), map_location=device) # load
ckpt = torch_load(attempt_download(path), map_location=device) # load
csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32
csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"]) # intersect
model.load_state_dict(csd, strict=False) # load