Fix overrides training cfg bug (ultralytics#10002)

Co-authored-by: UltralyticsAssistant <[email protected]> Co-authored-by: Glenn Jocher <[email protected]> Co-authored-by: Ultralytics AI Assistant <[email protected]> Co-authored-by: Laughing-q <[email protected]>
huilin66 · Apr 14, 2024 · 3bbbdcf · 3bbbdcf
1 parent a506a55
commit 3bbbdcf
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py
@@ -641,7 +641,12 @@ def train(
         checks.check_pip_update_available()
 
         overrides = yaml_load(checks.check_yaml(kwargs["cfg"])) if kwargs.get("cfg") else self.overrides
-        custom = {"data": DEFAULT_CFG_DICT["data"] or TASK2DATA[self.task]}  # method defaults
+        custom = {
+            # NOTE: handle the case when 'cfg' includes 'data'.
+            "data": overrides.get("data") or DEFAULT_CFG_DICT["data"] or TASK2DATA[self.task],
+            "model": self.overrides["model"],
+            "task": self.task,
+        }  # method defaults
         args = {**overrides, **custom, **kwargs, "mode": "train"}  # highest priority args on the right
         if args.get("resume"):
             args["resume"] = self.ckpt_path