From 70b466aa6727d14c0d2a05e708614cb854c41135 Mon Sep 17 00:00:00 2001
From: Dan Saunders <danjsaund@gmail.com>
Date: Mon, 31 Mar 2025 18:35:41 +0000
Subject: [PATCH] ray bugfix

---
 src/axolotl/train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/axolotl/train.py b/src/axolotl/train.py
index 56e378c7f..89f35d7eb 100644
--- a/src/axolotl/train.py
+++ b/src/axolotl/train.py
@@ -509,6 +509,7 @@ def train(
     # Save the trained model and cleanup
     save_trained_model(cfg, trainer, model, safe_serialization)
     create_model_card(cfg, trainer)
-    cleanup_distributed()
+    if not cfg.use_ray:
+        cleanup_distributed()
 
     return model, tokenizer, trainer