JLNTV_IPD/CogVideo
4
0
Fork 0
mirror of https://github.com/THUDM/CogVideo.git synced 2025-04-05 19:41:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add unload_model function

Browse Source
This commit is contained in:
OleehyO 2025-01-03 07:50:36 +00:00
parent a88c1ede69
commit f731c35f70
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
finetune/utils/memory_utils.py
View File

@ -51,6 +51,10 @@ def free_memory() -> None:
# TODO(aryan): handle non-cuda devices
def unload_model(model):
model.to("cpu")
def make_contiguous(x: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
if isinstance(x, torch.Tensor):
return x.contiguous()