From c62c5a5aad1394c9a4fa5863146bbb8ceb3e2bc6 Mon Sep 17 00:00:00 2001
From: JackCaoG <59073027+JackCaoG@users.noreply.github.com>
Date: Mon, 21 Nov 2022 10:38:56 -0800
Subject: [PATCH] Prevent multi register Lazy key (#4223)

* Prevent multi register Lazy key

* Update __init__.py
---
 torch_xla/__init__.py            | 4 +++-
 torch_xla/core/xrt_run_server.py | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
index c776918dd15..2bf2736be45 100644
--- 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
index 32673532ac7..ca629365772 100644
--- 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