Fix Ray train crashing after succeeding (#3542) [skip ci]

2026-03-25 07:38:28 -04:00
parent c2bd75aff6
commit 678ebb1bb2
1 changed files with 5 additions and 2 deletions
--- a/src/axolotl/cli/train.py
+++ b/src/axolotl/cli/train.py
@@ -84,8 +84,11 @@ def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
                storage_path=Path(parsed_cfg.output_dir).absolute().as_posix(),
            ),
        )
-        return trainer.fit()
-    return do_train(parsed_cfg, parsed_cli_args)
+
+        trainer.fit()
+        return
+
+    do_train(parsed_cfg, parsed_cli_args)


 def ray_train_func(kwargs: dict):