diff options
Diffstat (limited to 'modules/npu_specific.py')
| -rw-r--r-- | modules/npu_specific.py | 5 |
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() |