From f731c35f70e63d7fb32ce00b69b3277dc1b9403e Mon Sep 17 00:00:00 2001
From: OleehyO <leehy0357@gmail.com>
Date: Fri, 3 Jan 2025 07:50:36 +0000
Subject: [PATCH] Add unload_model function

---
 finetune/utils/memory_utils.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/finetune/utils/memory_utils.py b/finetune/utils/memory_utils.py
index a7a136b..b341d22 100644
--- a/finetune/utils/memory_utils.py
+++ b/finetune/utils/memory_utils.py
@@ -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()