From 88ca747e9ec39e35aacc19be5476d3220114b077 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Mon, 10 Feb 2025 19:46:19 +0100 Subject: [PATCH] nits --- llms/mlx_lm/tuner/grpo_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llms/mlx_lm/tuner/grpo_trainer.py b/llms/mlx_lm/tuner/grpo_trainer.py index ab1ab605..4a1e6bbf 100644 --- a/llms/mlx_lm/tuner/grpo_trainer.py +++ b/llms/mlx_lm/tuner/grpo_trainer.py @@ -159,7 +159,6 @@ def get_per_token_logps(model, inputs, lengths): logits = model(inputs).astype(mx.float16) logits = logits[:, :-1, :] targets = inputs[:, 1:] - mx.eval(logits) per_token_logps = [] for i in range(logits.shape[0]): seq_len = int(lengths[i]) - 1 @@ -172,6 +171,7 @@ def get_per_token_logps(model, inputs, lengths): axis=-1 ).squeeze(-1) per_token_logps.append(token_log_probs) + mx.eval(logits) return per_token_logps