diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 40979ec76f..0e9e28b48c 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -39,10 +39,9 @@ def PatchRL(FastLanguageModel):
     @contextmanager
     def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
         # Must use for_inference to allow inference in Unsloth
-        FastLanguageModel.for_inference(model)
-        with torch.inference_mode():
-            with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
-                yield unwrapped_model
+        with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
+            FastLanguageModel.for_inference(unwrapped_model)
+            yield unwrapped_model
         # Return back to training mode
         FastLanguageModel.for_training (model)
     pass