diff --git a/tests/test_scheduler_flax.py b/tests/test_scheduler_flax.py
index d29a8bfcc2..7928939f2d 100644
--- a/tests/test_scheduler_flax.py
+++ b/tests/test_scheduler_flax.py
@@ -859,7 +859,7 @@ def test_full_loop_no_noise(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 198.1542) < 1e-2
+            assert abs(result_sum - 198.1275) < 1e-2
             assert abs(result_mean - 0.2580) < 1e-3
         else:
             assert abs(result_sum - 198.1318) < 1e-2
@@ -872,8 +872,8 @@ def test_full_loop_with_set_alpha_to_one(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 185.4352) < 1e-2
-            assert abs(result_mean - 0.24145) < 1e-3
+            assert abs(result_sum - 186.83226) < 1e-2
+            assert abs(result_mean - 0.24327) < 1e-3
         else:
             assert abs(result_sum - 186.9466) < 1e-2
             assert abs(result_mean - 0.24342) < 1e-3
@@ -885,8 +885,8 @@ def test_full_loop_with_no_set_alpha_to_one(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 185.4352) < 1e-2
-            assert abs(result_mean - 0.2414) < 1e-3
+            assert abs(result_sum - 186.83226) < 1e-2
+            assert abs(result_mean - 0.24327) < 1e-3
         else:
             assert abs(result_sum - 186.9482) < 1e-2
             assert abs(result_mean - 0.2434) < 1e-3