make sure we flush printing the exception

pytorch · Feb 5, 2024 · b82e7c5 · b82e7c5
1 parent c95e362
commit b82e7c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
@@ -154,7 +154,7 @@ def _prepare_to_exit():
     if int(os.environ.get('PT_XLA_DEBUG', '0')):
       _summarize_fn_tracker()
   except Exception as e:
-    print('Caught an exception in the atexit callback: ', e)
+    print('Caught an exception in the atexit callback: ', e, flush=True)
     os._exit(1)