revert change to scheduler test

Author: entrpn

tests/schedulers/test_scheduler_flax.py

@@ -919,7 +919,7 @@ 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 - 186.94574) < 8e-2
+            assert abs(result_sum - 186.83226) < 8e-2
             assert abs(result_mean - 0.24327) < 1e-3
         else:
             assert abs(result_sum - 186.9466) < 1e-2
@@ -932,7 +932,7 @@ 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 - 186.94574) < 8e-2
+            assert abs(result_sum - 186.83226) < 8e-2
             assert abs(result_mean - 0.24327) < 1e-3
         else:
             assert abs(result_sum - 186.9482) < 1e-2