mirror of
https://github.com/RYDE-WORK/Langchain-Chatchat.git
synced 2026-01-22 23:11:05 +08:00
11 lines
303 B
Python
11 lines
303 B
Python
import torch.cuda
|
|
import torch.mps
|
|
import torch.backends
|
|
|
|
def torch_gc(DEVICE):
|
|
if torch.cuda.is_available():
|
|
with torch.cuda.device(DEVICE):
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
elif torch.backends.mps.is_available():
|
|
torch.mps.empty_cache() |