aboutsummaryrefslogtreecommitdiffstats
path: root/modules/npu_specific.py
diff options
context:
space:
mode:
authorwangshuai09 <391746016@qq.com>2024-01-31 02:46:53 +0000
committerwangshuai09 <391746016@qq.com>2024-01-31 04:29:58 +0000
commitcc3f604310458eed7d26456c1b3934d582283ffe (patch)
tree44cbaa680fd3afeea5799fa75af4d52872f646ba /modules/npu_specific.py
parent74ff85a1a1ee4cce432b1c7d33c1eda831f68d48 (diff)
downloadstable-diffusion-webui-gfx803-cc3f604310458eed7d26456c1b3934d582283ffe.tar.gz
stable-diffusion-webui-gfx803-cc3f604310458eed7d26456c1b3934d582283ffe.tar.bz2
stable-diffusion-webui-gfx803-cc3f604310458eed7d26456c1b3934d582283ffe.zip
Update
Diffstat (limited to 'modules/npu_specific.py')
-rw-r--r--modules/npu_specific.py5
1 files changed, 1 insertions, 4 deletions
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()