From f87b447b2d1ed13a4a101c7ca5d70e6d1beab2fb Mon Sep 17 00:00:00 2001 From: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Date: Mon, 16 Dec 2024 18:59:05 +0800 Subject: [PATCH] Fix `model.save()` for model YAMLs (#18212) Signed-off-by: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Signed-off-by: Glenn Jocher Co-authored-by: Glenn Jocher --- ultralytics/engine/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py index db8d87eb..8affd958 100644 --- a/ultralytics/engine/model.py +++ b/ultralytics/engine/model.py @@ -115,7 +115,7 @@ class Model(nn.Module): self.predictor = None # reuse predictor self.model = None # model object self.trainer = None # trainer object - self.ckpt = None # if loaded from *.pt + self.ckpt = {} # if loaded from *.pt self.cfg = None # if loaded from *.yaml self.ckpt_path = None self.overrides = {} # overrides for trainer object @@ -807,7 +807,7 @@ class Model(nn.Module): # Update model and cfg after training if RANK in {-1, 0}: ckpt = self.trainer.best if self.trainer.best.exists() else self.trainer.last - self.model, _ = attempt_load_one_weight(ckpt) + self.model, self.ckpt = attempt_load_one_weight(ckpt) self.overrides = self.model.args self.metrics = getattr(self.trainer.validator, "metrics", None) # TODO: no metrics returned by DDP return self.metrics