diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py
index c52d14f402..3478dfc31a 100644
--- a/unsloth/models/__init__.py
+++ b/unsloth/models/__init__.py
@@ -20,3 +20,4 @@ from .mistral import FastMistralModel
 from .qwen2   import FastQwen2Model
 from .dpo     import PatchDPOTrainer, PatchKTOTrainer
 from ._utils  import is_bfloat16_supported
+from .rl      import PatchRL
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index efe2d33e01..2aa8f02659 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -13,14 +13,15 @@
 # limitations under the License.
 
 __all__ = [
-    "patch_rl",
+    "PatchRL",
 ]
 
-from trl.models.utils import unwrap_model_for_generation
-from contextlib import contextmanager
 
+def PatchRL(FastLanguageModel):
 
-def patch_rl(FastLanguageModel):
+    from trl.models.utils import unwrap_model_for_generation
+    from contextlib import contextmanager
+    
     @contextmanager
     def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
         FastLanguageModel.for_inference(model)