weiweiw
/
Langchain-Chatchat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update utils

This commit is contained in:
imClumsyPanda 2023-05-03 10:08:35 +08:00
parent 4c8429cb15
commit e567b862cc
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
utils/__init__.py
View File

@ -1,5 +1,4 @@
import torch.cuda import torch.cuda
import torch.mps
import torch.backends import torch.backends
def torch_gc(DEVICE): def torch_gc(DEVICE):
@ -8,4 +7,9 @@ def torch_gc(DEVICE):
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
elif torch.backends.mps.is_available(): elif torch.backends.mps.is_available():
try:
import torch.mps
torch.mps.empty_cache() torch.mps.empty_cache()
except Exception as e:
print(e)
print("如果您使用的是 macOS 建议将 pytorch 版本升级至 2.0.0 或更高版本，以支持及时清理 torch 产生的内存占用。")