From 2822acc98aa5db82977edd2a722c5ad4d3698087 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Thu, 24 Oct 2024 06:19:19 -0700
Subject: [PATCH] fix synch

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

diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py
index e9388dbe..a7a998f4 100644
--- a/llms/mlx_lm/utils.py
+++ b/llms/mlx_lm/utils.py
@@ -58,8 +58,7 @@ def wired_limit(model):
     try:
         yield None
     finally:
-        # TODO... expose a synchronize??
-        mx.zeros((1,)).item()
+        mx.synchronize()
         mx.metal.set_wired_limit(old_limit)