From 5cf02ec21ea71de485bcc6b83f66b521eb7f79b5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 5 Feb 2025 05:19:37 -0800 Subject: [PATCH] Update rl.py --- unsloth/models/rl.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 2aa8f02659..2bd602e099 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -21,11 +21,12 @@ def PatchRL(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) - yield unwrap_model_for_generation(model, *args, **kwargs) + with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model: + yield unwrapped_model FastLanguageModel.for_training (model) pass