diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 281e3a71df..6584e2a426 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -510,7 +510,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" # logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred - return logits, entropies # logps, entropies + return logits.detach(), entropies # logps, entropies # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770