From 06811d3826966ecf4e423681d974a70ab3a48de8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 6 Feb 2025 02:31:01 -0800
Subject: [PATCH] Update rl.py

---
 unsloth/models/rl.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index c785870306..22e1e0f6cb 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -353,6 +353,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
             source,
         )
 
+        # .state_dict()
+        source = re.sub(
+            r"\.state_dict\(\)",
+            r"",
+            source,
+        )
+        
         # Replace self.llm.generate and self.llm.chat
         lora_name = trainer_file + "_lora_model"
         source = re.sub(
@@ -382,6 +389,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         f"class {RLTrainer_name}", f"class Unsloth{RLTrainer_name}", 1
     )
 
+    # Create new class in compiled cache and import it
     module = create_new_function(
         RLTrainer_name,
         RLTrainer_source,