fix #124

fudan-generative-vision · Jul 3, 2024 · 9fc2817 · 9fc2817
1 parent c1ce01f
commit 9fc2817
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/scripts/train_stage1.py b/scripts/train_stage1.py
@@ -16,6 +16,7 @@
 """
 
 import argparse
+import copy
 import logging
 import math
 import os
@@ -211,6 +212,7 @@ def log_validation(
     logger.info("Running validation... ")
 
     ori_net = accelerator.unwrap_model(net)
+    ori_net = copy.deepcopy(ori_net)
     reference_unet = ori_net.reference_unet
     denoising_unet = ori_net.denoising_unet
     face_locator = ori_net.face_locator
@@ -278,6 +280,7 @@ def log_validation(
         canvas.save(out_file)
 
     del pipe
+    del ori_net
     torch.cuda.empty_cache()
 
     return pil_images