From a0ce0594f61e2b6e7003d88623b66704cdf49b07 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Thu, 29 Aug 2024 10:33:43 -0700
Subject: [PATCH] Temporarily remove async_eval

---
 llms/mlx_lm/generate.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llms/mlx_lm/generate.py b/llms/mlx_lm/generate.py
index aa7a4a2f..245ad155 100644
--- a/llms/mlx_lm/generate.py
+++ b/llms/mlx_lm/generate.py
@@ -249,8 +249,10 @@ def main():
         draft_model=draft_model,
         num_draft_tokens=args.num_draft_tokens,
     )
-    if not args.verbose:
+
+    if not args.verbose and mx.distributed.init().rank() == 0:
         print(response)
+    mx.synchronize()
 
 
 if __name__ == "__main__":