diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index e5101662b8..515c6587f7 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -47,8 +47,8 @@ def PatchRL(FastLanguageModel):
     from contextlib import contextmanager
 
     @contextmanager
-    def unsloth_unwrap_model_for_generation(model, accelerator):
-        with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
+    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)