From e34ecb79b420b7ed1125c225f54fee5656b07dbf Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Fri, 13 Dec 2024 18:44:56 -0800
Subject: [PATCH] fix

---
 llms/mlx_lm/utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py
index 102512c1..cc420bf0 100644
--- a/llms/mlx_lm/utils.py
+++ b/llms/mlx_lm/utils.py
@@ -260,6 +260,8 @@ def generate_step(
         kv_bits=kv_bits,
     )
 
+    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
+
     def _step(y):
         with mx.stream(generation_stream):
             logits = model(y[None], cache=prompt_cache)