fix(tests/sampling): use float32 for testing

ml-gde · Oct 9, 2024 · 55f378b · 55f378b
1 parent 12dd0da
commit 55f378b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
@@ -22,14 +22,14 @@ def test_get_noise(self):
             num_samples=1,
             height=height,
             width=width,
-            dtype=jax.dtypes.bfloat16,
+            dtype=jax.numpy.float32,
             seed=jax.random.PRNGKey(seed=42),
         )
         x_torch = torch_get_noise(
             num_samples=1,
             height=height,
             width=width,
-            dtype=torch.bfloat16,
+            dtype=torch.float32,
             seed=42,
             device="cuda" if torch.cuda.is_available() else "cpu",
         )