Added effects_barrier to fix test issue with JAX 0.4.33+

patrick-kidger · Oct 21, 2024 · cbf04ff · cbf04ff
1 parent d80ff7a
commit cbf04ff
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/test/test_progress_meter.py b/test/test_progress_meter.py
@@ -35,6 +35,7 @@ def solve(t0):
     for num_lines, solve_fn in solve_fns:
         capfd.readouterr()
         solve_fn()
+        jax.effects_barrier()
         captured = capfd.readouterr()
         err = captured.err.strip()
         assert re.match("0.00%|[ ]+|", err.split("\r", 1)[0])