MLU memory optimization

This commit is contained in:
huzhan 2025-04-02 16:02:13 +08:00
parent ab5413351e
commit d3a2ee5ced

View File

@ -1237,6 +1237,8 @@ def soft_empty_cache(force=False):
torch.xpu.empty_cache()
elif is_ascend_npu():
torch.npu.empty_cache()
elif is_mlu():
torch.mlu.empty_cache()
elif torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()