theluyuan/RWKV-Runner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

prevent 'torch' has no attribute 'cuda' error in torch_gc, so user can use CPU or WebGPU (#302)

Browse Source
This commit is contained in:
josc146 2024-03-11 19:04:19 +08:00
parent dd89041f72
commit bb1a6191b0
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
backend-python/utils/torch.py
View File

@ -19,9 +19,12 @@ def set_torch():
def torch_gc(): def torch_gc():
import torch try:
import torch
if torch.cuda.is_available(): if torch.cuda.is_available():
with torch.cuda.device(0): with torch.cuda.device(0):
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
except:
pass # prevent 'torch' has no attribute 'cuda' error, so user can use CPU or WebGPU