From cc3f604310458eed7d26456c1b3934d582283ffe Mon Sep 17 00:00:00 2001
From: wangshuai09 <391746016@qq.com>
Date: Wed, 31 Jan 2024 10:46:53 +0800
Subject: Update

---
 modules/initialize.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

(limited to 'modules/initialize.py')

diff --git a/modules/initialize.py b/modules/initialize.py
index cc34fd6f..f7313ff4 100644
--- a/modules/initialize.py
+++ b/modules/initialize.py
@@ -143,10 +143,7 @@ def initialize_rest(*, reload_script_modules=False):
         by that time, so we apply optimization again.
         """
         from modules import devices
-        # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
-        if devices.npu_specific.has_npu:
-            import torch
-            torch.npu.set_device(0)
+        devices.torch_npu_set_device()
 
         shared.sd_model  # noqa: B018
 
-- 
cgit v1.2.3