diff options
| author | wangshuai09 <391746016@qq.com> | 2024-01-31 02:46:53 +0000 |
|---|---|---|
| committer | wangshuai09 <391746016@qq.com> | 2024-01-31 04:29:58 +0000 |
| commit | cc3f604310458eed7d26456c1b3934d582283ffe (patch) | |
| tree | 44cbaa680fd3afeea5799fa75af4d52872f646ba /modules/devices.py | |
| parent | 74ff85a1a1ee4cce432b1c7d33c1eda831f68d48 (diff) | |
| download | stable-diffusion-webui-gfx803-cc3f604310458eed7d26456c1b3934d582283ffe.tar.gz stable-diffusion-webui-gfx803-cc3f604310458eed7d26456c1b3934d582283ffe.tar.bz2 stable-diffusion-webui-gfx803-cc3f604310458eed7d26456c1b3934d582283ffe.zip | |
Update
Diffstat (limited to 'modules/devices.py')
| -rw-r--r-- | modules/devices.py | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/modules/devices.py b/modules/devices.py index c737162a..28c0c54d 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -88,9 +88,16 @@ def torch_gc(): xpu_specific.torch_xpu_gc() if npu_specific.has_npu: + torch_npu_set_device() npu_specific.torch_npu_gc() +def torch_npu_set_device(): + # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue + if npu_specific.has_npu: + torch.npu.set_device(0) + + def enable_tf32(): if torch.cuda.is_available(): |