Merge pull request #23870 from Zantares:tenglu/flush_output

PiperOrigin-RevId: 679639244
jax-ml · Sep 27, 2024 · df042fd · df042fd
2 parents a616758 + a31f79c
commit df042fd
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/api_test.py b/tests/api_test.py
@@ -10906,6 +10906,8 @@ def test_call_wrapped_second_phase_cleanup(self):
 class EnvironmentInfoTest(jtu.JaxTestCase):
   @parameterized.parameters([True, False])
   def test_print_environment_info(self, return_string):
+    # Flush stdout buffer before checking.
+    sys.stdout.flush()
     with jtu.capture_stdout() as stdout:
       result = jax.print_environment_info(return_string=return_string)
     if return_string: