From 3ebb7ebaf834e54129bee4d98b45b8e2895682e6 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 5 Feb 2025 06:48:54 -0800
Subject: [PATCH] Update rl.py

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

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 6603346fd0..d77a4b3784 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -39,15 +39,14 @@ def PatchRL(FastLanguageModel):
     @contextmanager
     def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
         # Must use for_inference to allow inference in Unsloth
-        print("$$$$$$$$$$$$$$")
-        FastLanguageModel.for_inference(model)
-        try:
-            with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
+        with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
+            FastLanguageModel.for_inference(unwrapped_model)
+            try:
                 yield unwrapped_model
-        finally:
-            # Finally return back training
-            FastLanguageModel.for_training(model)
-            print("###############")
+            finally:
+                # Finally return back training
+                FastLanguageModel.for_training(model)
+            pass
         pass
     pass