From 52ae65ac3bc414157a04a5706ade18e3e1b6a283 Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Thu, 1 Jun 2023 04:32:16 +0530 Subject: [PATCH] Optionally `train_args` on attempt_load_weights() (#2928) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- ultralytics/nn/tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index aaeef5b..7af3f34 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -532,7 +532,7 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=False): ensemble = Ensemble() for w in weights if isinstance(weights, list) else [weights]: ckpt, w = torch_safe_load(w) # load ckpt - args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} # combine model and default args, preferring model args + args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} if 'train_args' in ckpt else None # combined args model = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model # Model compatibility updates