diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 993173a53..de3e85c08 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -50,7 +50,6 @@ parser.add_argument("--input-directory", type=str, default=None, help="Set the C
 parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
 parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
 parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use. All other devices will not be visible.")
-parser.add_argument("--ascend-device", type=int, default=None, metavar="ASCEND_DEVICE_ID", help="Set the id of the ascend device this instance will use. All other devices will not be visible.")
 parser.add_argument("--default-device", type=int, default=None, metavar="DEFAULT_DEVICE_ID", help="Set the id of the default device, all other devices will stay visible.")
 cm_group = parser.add_mutually_exclusive_group()
 cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
diff --git a/main.py b/main.py
index 0ab33d603..a6eec83af 100644
--- a/main.py
+++ b/main.py
@@ -103,18 +103,23 @@ def execute_prestartup_script():
 apply_custom_paths()
 execute_prestartup_script()
 
+if 'torch' in sys.modules:
+    logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
 
 # Main code
 import asyncio
 import shutil
 import threading
 import gc
-
+import comfy.model_management
 
 if os.name == "nt":
     os.environ['MIMALLOC_PURGE_DELAY'] = '0'
 
 if __name__ == "__main__":
+    device = comfy.model_management.get_torch_device()
+    device_name = comfy.model_management.get_torch_device_name(device)
+
     if args.default_device is not None:
         default_dev = args.default_device
         devices = list(range(32))
@@ -125,12 +130,13 @@ if __name__ == "__main__":
         os.environ['HIP_VISIBLE_DEVICES'] = str(devices)
 
     if args.cuda_device is not None:
-        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
-        os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device)
-        logging.info("Set cuda device to: {}".format(args.cuda_device))
-    elif args.ascend_device is not None:
-        os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.ascend_device)
-        logging.info("Set npu device to: {}".format(args.ascend_device))
+        if device_name != 'npu':
+            os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
+            os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device)
+            logging.info("Set cuda device to: {}".format(args.cuda_device))
+        else:
+            os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device)
+            logging.info("Set npu device to: {}".format(args.cuda_device))
 
     if args.oneapi_device_selector is not None:
         os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector
@@ -142,8 +148,7 @@ if __name__ == "__main__":
 
     import cuda_malloc
 
-if 'torch' in sys.modules:
-    logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
+
 
 import comfy.utils
 
@@ -151,7 +156,6 @@ import execution
 import server
 from protocol import BinaryEventTypes
 import nodes
-import comfy.model_management
 import comfyui_version
 import app.logger
 import hook_breaker_ac10a0