Fixes for CUDA CI (#3413)

python/tests/test_conv.py

@@ -1,6 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import math
+import os
 import unittest
 from itertools import permutations
 
@@ -1151,6 +1152,7 @@ class TestConv(mlx_tests.MLXTestCase):
             )
             self.assertEqual(grads.shape, k_shape)
 
+    @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "flaky in CI")
     def test_conv_1d_with_2d(self):
         x = mx.random.uniform(shape=(2, 10, 16))
         y = mx.random.normal(shape=(16, 3, 16))

python/tests/test_fast_sdpa.py

@@ -1,4 +1,5 @@
 import math
+import os
 import unittest
 from itertools import product
 
@@ -443,6 +444,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
                 self.assertLessEqual(mx.max(diff).item(), atol)
 
     @unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU")
+    @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "not enough memory")
     def test_sdpa_long_masked_sequence(self):
         # Test for int16 overflow in steel_attention_nax.h mask
         # indexing (col_pos declared as short, overflows when kL > 32767).