diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
index 3337a00..18633cb 100644
--- a/.github/workflows/pull_request.yml
+++ b/.github/workflows/pull_request.yml
@@ -40,5 +40,5 @@ jobs:
         run: |
           curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
           unzip test_data.zip
-          HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
+          METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
           mlx.launch -n 2 tests/model_parallel_tests.py
diff --git a/tests/test_losses.py b/tests/test_losses.py
index 3436b76..e9015a8 100644
--- a/tests/test_losses.py
+++ b/tests/test_losses.py
@@ -13,21 +13,21 @@ class TestLosses(unittest.TestCase):
         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)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
 
         with mx.stream(mx.cpu):
             expected = kl_div_loss(logits_q, logits_p)
         kl = kl_div_loss(logits_q, logits_p)
 
-        self.assertTrue(mx.allclose(kl, expected, rtol=1e-4))
+        self.assertTrue(mx.allclose(kl, expected))
 
     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)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
 
         with mx.stream(mx.cpu):
             expected = js_div_loss(logits_q, logits_p)
@@ -39,9 +39,9 @@ class TestLosses(unittest.TestCase):
         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)
-        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
+        cotan = mx.random.normal((2, 4))
 
         with mx.stream(mx.cpu):
             expected = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]
@@ -53,9 +53,9 @@ class TestLosses(unittest.TestCase):
         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)
-        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
+        cotan = mx.random.normal((2, 4))
 
         with mx.stream(mx.cpu):
             expected = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]