diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index d77a4b3784..3129488f32 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -39,7 +39,7 @@ def PatchRL(FastLanguageModel):
     @contextmanager
     def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
         # Must use for_inference to allow inference in Unsloth
-        with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
+        with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
             FastLanguageModel.for_inference(unwrapped_model)
             try:
                 yield unwrapped_model