diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py
index b37a13d..7c3307d 100644
--- a/mlx_lm/generate.py
+++ b/mlx_lm/generate.py
@@ -1228,7 +1228,7 @@ class BatchGenerator:
             batch.tokens,
         )
 
-        mx.async_eval(batch.y, batch.logprobs)
+        mx.async_eval(batch.y, batch.logprobs, batch.tokens)
 
         y = y.tolist()
         toc = time.perf_counter()