From 5c2c53afee07b5a5c8d09bd66ae5d3d550304910 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 24 Nov 2025 22:13:17 -0800
Subject: [PATCH] Update rl.py

---
 unsloth/models/rl.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index bfbd0bf103..1a06277d85 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -229,6 +229,7 @@ from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling
 from transformers.training_args import ParallelMode
 
 # Wrap trainer with padding to right and enable training mode
+# Also patches W&B since multiple runs must use wandb.finish()
 import functools
 from types import MethodType
 def prepare_for_training_mode(f):
@@ -241,6 +242,12 @@ def prepare_for_training_mode(f):
         # Return inference mode
         if hasattr(self, 'model') and hasattr(self.model, "for_inference"):
             self.model.for_inference()
+        # Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run
+        try:
+            import wandb
+            wandb.finish()
+        except:
+            pass
         return output
     return wrapper
 pass