zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-12-15 17:58:54 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

smoll fix

Browse Source
This commit is contained in:
Goekdeniz-Guelmez
2025-02-26 15:21:57 +01:00
parent ef6ff92add
commit fab2dc2688
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
llms/mlx_lm/tuner/grpo_trainer.py
View File

@@ -279,6 +279,7 @@ def grpo_loss(
reward_weights = mx.array(reward_weights, dtype=mx.float32)
else:
reward_weights = mx.ones(len(reward_funcs), dtype=mx.float32)
rewards = (rewards * mx.expand_dims(reward_weights, 0)).sum(axis=1)
# Reshape rewards and compute advantages