diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 932a29f78a..2282e8b313 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -43,8 +43,7 @@ def PatchRL(FastLanguageModel):
             FastLanguageModel.for_inference(unwrapped_model)
             yield unwrapped_model
         # Return back to training mode
-        FastLanguageModel.for_training (model)
-        yield model
+        FastLanguageModel.for_training(model)
     pass
 
     import trl.trainer