From 1c9ae1eaa1e16a2a159e2687e22dc172351b9358 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Thu, 2 Oct 2025 15:40:04 -0700
Subject: [PATCH] cuda fix flaky test (#2646)

---
 python/tests/test_array.py    | 1 +
 python/tests/test_autograd.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/python/tests/test_array.py b/python/tests/test_array.py
index e932382b1..a8f9474ef 100644
--- a/python/tests/test_array.py
+++ b/python/tests/test_array.py
@@ -2005,6 +2005,7 @@ class TestArray(mlx_tests.MLXTestCase):
             b = mx.reshape(b, [])
             return b
 
+        mx.synchronize()
         t()
         gc.collect()
         expected = get_mem()
diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py
index 5722071f6..38bb6089d 100644
--- a/python/tests/test_autograd.py
+++ b/python/tests/test_autograd.py
@@ -746,6 +746,7 @@ class TestAutograd(mlx_tests.MLXTestCase):
             mx.checkpoint,
         ]:
             mx.synchronize()
+            gc.collect()
             mem_pre = mx.get_active_memory()
 
             def outer():