Re-enable async dispatch on JAX CPU in equations_test.py

The tsan error has been fixed. PiperOrigin-RevId: 660183448
google · Aug 7, 2024 · 48f69f5 · 48f69f5
1 parent 2d43ad9
commit 48f69f5
Showing 1 changed file with 0 additions and 3 deletions.
diff --git a/jax_cfd/ml/equations_test.py b/jax_cfd/ml/equations_test.py
@@ -293,7 +293,4 @@ def test_epd_modules(
 
 
 if __name__ == '__main__':
-  # Temporarily disable async dispatch on JAX CPU due to tsan error.
-  jax.config.update('jax_cpu_enable_async_dispatch', False)
-
   absltest.main()