From 290cfa1122b779b9acc0ec556d99c70be715e824 Mon Sep 17 00:00:00 2001
From: Yu Li <lyo.gavin@gmail.com>
Date: Fri, 1 Dec 2023 17:50:34 -0600
Subject: [PATCH] make sure its on cuda

---
 air_llm/airllm/airllm.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/air_llm/airllm/airllm.py b/air_llm/airllm/airllm.py
index a414a3a..9f10469 100644
--- a/air_llm/airllm/airllm.py
+++ b/air_llm/airllm/airllm.py
@@ -209,9 +209,10 @@ def split_and_save_layers(checkpoint_path, layer_shards_saving_path=None, splitt
 
         # Free memory
         for k in layer_state_dict.keys():
-            del state_dict[k]
+            if k in state_dict:
+                del state_dict[k]
         del layer_state_dict
-        gc.collect()
+        clean_memory()
 
     return str(saving_path)