提交 8db6cd51 作者: imClumsyPanda

update utils

上级 a0b312d7
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): ...@@ -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():
torch.mps.empty_cache() try:
\ No newline at end of file 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 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论