From 894ea97496d65d2cc4fd6f5bd6df642d4e6bf388 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 5 Mar 2025 23:17:51 -0800
Subject: [PATCH] Update rl.py

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

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index c9ea922272..25555e2628 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -284,6 +284,18 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         extra_args += eval_changes
     pass
 
+    # Force logits to be produced if preprocess_logits_for_metrics or compute_metrics is used
+    if "model" in call_args:
+        logits_check = \
+        "_output_logits = False"\
+        "if locals().get('compute_metrics', None) is not None: _output_logits = True\n"\
+        "if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True\n"\
+        "if _output_logits:\n"\
+        "    import os\n"\
+        "    os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"
+        extra_args += logits_check
+    pass
+
     # Check max_seq_length
     if "model" in call_args:
         length_check = \