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

+4 -2
View File
@@ -44,6 +44,8 @@ if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
from ultralytics.utils.patches import torch_load
import val as validate # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
@@ -212,7 +214,7 @@ def train(hyp, opt, device, callbacks):
if pretrained:
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak
ckpt = torch_load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak
model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create
exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys
csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32
@@ -646,7 +648,7 @@ def main(opt, callbacks=Callbacks()):
with open(opt_yaml, errors="ignore") as f:
d = yaml.safe_load(f)
else:
d = torch.load(last, map_location="cpu")["opt"]
d = torch_load(last, map_location="cpu")["opt"]
opt = argparse.Namespace(**d) # replace
opt.cfg, opt.weights, opt.resume = "", str(last), True # reinstate
if is_url(opt_data):