2023-05-03 13:14:36 +08:00
|
|
|
import torch
|
2023-05-01 23:51:29 +08:00
|
|
|
|
|
|
|
|
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():
|
2023-05-03 10:08:35 +08:00
|
|
|
try:
|
|
|
|
|
torch.mps.empty_cache()
|
|
|
|
|
except Exception as e:
|
|
|
|
|
print(e)
|
|
|
|
|
print("如果您使用的是 macOS 建议将 pytorch 版本升级至 2.0.0 或更高版本，以支持及时清理 torch 产生的内存占用。")
|
