cuda fix flaky test (#2646)

2025-10-17 23:08:11 +08:00 · 2025-10-02 15:40:04 -07:00
parent c2c3e0b0a2
commit 1c9ae1eaa1
2 changed files with 2 additions and 0 deletions
--- 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()
--- 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():