oobabooga/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2026-04-06 15:13:38 +00:00
Code Issues Projects Releases Packages Wiki Activity

Add Ascend NPU support (basic) (#5541)

Browse source
This commit is contained in:
wangshuai09 2024-04-12 05:42:20 +08:00 • committed by GitHub
parent a90509d82e
commit fd4e46bce2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 35 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

4
modules/callbacks.py
View file

@ -5,7 +5,7 @@ from threading import Thread
import torch
import transformers
from transformers import is_torch_xpu_available
from transformers import is_torch_npu_available, is_torch_xpu_available
import modules.shared as shared
@ -99,5 +99,7 @@ def clear_torch_cache():
if not shared.args.cpu:
if is_torch_xpu_available():
torch.xpu.empty_cache()
elif is_torch_npu_available():
torch.npu.empty_cache()
else:
torch.cuda.empty_cache()