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/npu_specific.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

(limited to 'modules/npu_specific.py')

diff --git a/modules/npu_specific.py b/modules/npu_specific.py
index d8aebf9c..94100691 100644
--- a/modules/npu_specific.py
+++ b/modules/npu_specific.py
@@ -8,11 +8,10 @@ def check_for_npu():
     if importlib.util.find_spec("torch_npu") is None:
         return False
     import torch_npu
-    torch_npu.npu.set_device(0)
 
     try:
         # Will raise a RuntimeError if no NPU is found
-        _ = torch.npu.device_count()
+        _ = torch_npu.npu.device_count()
         return torch.npu.is_available()
     except RuntimeError:
         return False
@@ -25,8 +24,6 @@ def get_npu_device_string():
 
 
 def torch_npu_gc():
-    # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
-    torch.npu.set_device(0)
     with torch.npu.device(get_npu_device_string()):
         torch.npu.empty_cache()
 
-- 
cgit v1.2.3