Skip to content

GitLab

Unverified Commit 2e980ac9 authored by Anton Lozhkov's avatar Anton Lozhkov Committed by GitHub
Browse files

[Tests] Adjust TPU test values (#1233) 

* [Tests] Adjust TPU test values

* slow tests

* remaining refs
parent 0feb21a1
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 5 deletions
tests/test_scheduler_flax.py
View file @ 2e980ac9
......@@ -859,7 +859,7 @@ class FlaxPNDMSchedulerTest(FlaxSchedulerCommonTest):
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 @@ class FlaxPNDMSchedulerTest(FlaxSchedulerCommonTest):
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 @@ class FlaxPNDMSchedulerTest(FlaxSchedulerCommonTest):
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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment