diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py
index 4c979e02dcd90..a3637cc890c0d 100644
--- a/python/taichi/lang/__init__.py
+++ b/python/taichi/lang/__init__.py
@@ -237,12 +237,8 @@ def benchmark(func, repeat=100, args=()):
     return elapsed / repeat
 
 
-wanted_archs = None
-
-
 def set_wanted_archs(archs):
-    global wanted_archs
-    wanted_archs = archs
+    os.environ['TI_WANTED_ARCHS'] = ','.join(archs)
 
 
 def supported_archs():
@@ -254,7 +250,8 @@ def supported_archs():
         archs.append(metal)
     if ti.core.with_opengl():
         archs.append(opengl)
-    if wanted_archs is not None:
+    wanted_archs = os.environ.get('TI_WANTED_ARCHS', '').split(',')
+    if len(wanted_archs):
         archs, old_archs = [], archs
         for arch in old_archs:
             if ti.core.arch_name(arch) in wanted_archs:
diff --git a/python/taichi/main.py b/python/taichi/main.py
index 6193662bbea27..1ff29a2522112 100644
--- a/python/taichi/main.py
+++ b/python/taichi/main.py
@@ -109,6 +109,7 @@ def main(debug=False):
     print()
     import taichi as ti
     if args.arch is not None:
+        print(f'Running on Arch={args.arch}')
         ti.set_wanted_archs(args.arch.split(','))
 
     if mode == 'help':