Prevent multi register Lazy key (#4223)

* Prevent multi register Lazy key * Update __init__.py
pytorch · Nov 21, 2022 · c62c5a5 · c62c5a5
1 parent 2209ed5
commit c62c5a5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
@@ -129,4 +129,6 @@ def _init_xla_lazy_backend():
 
 atexit.register(_prepare_to_exit)
 _apply_patches()
-_init_xla_lazy_backend()
+# _init_xla_lazy_backend should not be called more than once.
+if os.environ.get('INIT_LAZY_BACKEND', '1') != '0':
+  _init_xla_lazy_backend()
diff --git a/torch_xla/core/xrt_run_server.py b/torch_xla/core/xrt_run_server.py
@@ -6,6 +6,7 @@
 import sys
 
 from pathlib import Path
+os.environ['INIT_LAZY_BACKEND'] = '0'
 from torch_xla.__init__ import server_is_alive, XRT_RUN_SERVER_PROCESS, XRT_SERVER_REGEX