From 678ebb1bb2a6edb3230d6c873e32f06e6ee1b754 Mon Sep 17 00:00:00 2001
From: Matthew Hambrecht <14303543+mhambre@users.noreply.github.com>
Date: Wed, 25 Mar 2026 07:38:28 -0400
Subject: [PATCH] Fix Ray train crashing after succeeding (#3542) [skip ci]

---
 src/axolotl/cli/train.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/axolotl/cli/train.py b/src/axolotl/cli/train.py
index 6b3bfbd57..2d472ed10 100644
--- 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):