From f940cf3a952659f4b9e9c3d0d69ea5b083a533d8 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Tue, 2 Dec 2025 13:33:11 -0800 Subject: [PATCH] fix flaky test (#643) --- tests/test_losses.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_losses.py b/tests/test_losses.py index 8ac5dca..3436b76 100644 --- a/tests/test_losses.py +++ b/tests/test_losses.py @@ -11,6 +11,7 @@ class TestLosses(unittest.TestCase): def test_kl_div_loss(self): self.assertTrue(can_run_metal()) + mx.random.seed(0) logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) @@ -23,6 +24,7 @@ class TestLosses(unittest.TestCase): def test_js_div_loss(self): self.assertTrue(can_run_metal()) + mx.random.seed(0) logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) @@ -35,6 +37,7 @@ class TestLosses(unittest.TestCase): def test_kl_div_loss_vjp(self): self.assertTrue(can_run_metal()) + mx.random.seed(0) logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) @@ -48,6 +51,7 @@ class TestLosses(unittest.TestCase): def test_js_div_loss_vjp(self): self.assertTrue(can_run_metal()) + mx.random.seed(0) logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32) logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)