提交 e567b862 作者: imClumsyPanda

update utils

上级 4c8429cb
import torch.cuda
import torch.mps
import torch.backends
def torch_gc(DEVICE):
......@@ -8,4 +7,9 @@ def torch_gc(DEVICE):
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
elif torch.backends.mps.is_available():
try:
import torch.mps
torch.mps.empty_cache()
except Exception as e:
print(e)
print("如果您使用的是 macOS 建议将 pytorch 版本升级至 2.0.0 或更高版本,以支持及时清理 torch 产生的内存占用。")
\ No newline at end of file
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论