add tokens to eval to avoid large graphs when they are not used (#924)

2026-02-23 14:38:08 -08:00
parent 65725dcec2
commit 720f2369ba
1 changed files with 1 additions and 1 deletions
@@ -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()