diff options
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(): |