diff --git a/facechain/merge_lora.py b/facechain/merge_lora.py
index dbd19833..82690df0 100644
--- a/facechain/merge_lora.py
+++ b/facechain/merge_lora.py
@@ -54,7 +54,7 @@ def merge_lora(pipeline,
             curr_layer = pipeline.unet
 
         temp_name = layer_infos.pop(0)
-        while len(layer_infos) > -1:
+        while True:
             try:
                 curr_layer = curr_layer.__getattr__(temp_name)
                 if len(layer_infos) > 0:
@@ -138,7 +138,7 @@ def restore_lora(pipeline,
             curr_layer = pipeline.unet
 
         temp_name = layer_infos.pop(0)
-        while len(layer_infos) > -1:
+        while True:
             try:
                 curr_layer = curr_layer.__getattr__(temp_name)
                 if len(layer_infos) > 0:
@@ -162,5 +162,6 @@ def restore_lora(pipeline,
 
         curr_layer.weight.data = curr_layer.weight.data.to(device)
         curr_layer.weight.data = curr_layer.weight.data_restore.clone()
+        del curr_layer.weight.data_restore
 
     return pipeline
\ No newline at end of file