ordering

ml-explore · Dec 19, 2024 · e34273e · e34273e
1 parent de23433
commit e34273e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llms/mlx_lm/evaluate.py b/llms/mlx_lm/evaluate.py
@@ -209,10 +209,10 @@ def loglikelihood(self, requests) -> list[tuple[float, bool]]:
         if group is not None:
             per_group = int(np.ceil(num_results / group.size()))
             scores = mx.pad(scores, ((0, per_group - len(scores)),))
+            is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
             scores = mx.distributed.all_gather(scores[mx.newaxis], stream=mx.cpu)
-            scores = scores.T.reshape(-1)
             is_greedy = mx.distributed.all_gather(is_greedy[mx.newaxis], stream=mx.cpu)
-            is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
+            scores = scores.T.reshape(-1)
             is_greedy = is_greedy.T.reshape(-1)
 
         scores = np.array(scores[:num_results])