diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 3a9d651d11..c9ea922272 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -60,7 +60,7 @@ def PatchRL(FastLanguageModel):
     def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
         with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
             # Put the model in inference mode.
-            FastLanguageModel.for_inference(unwrapped_model)
+            FastLanguageModel.for_inference(model)
 
             # We must use .clone for Unsloth since we force inference_mode
             # Rather we should have used no_grad