Merge pull request THUDM#383 from LittleNyima/fix-deprecation

fix deprecation of clear_objs_and_retain_memory
wenc-k · Oct 1, 2024 · a9a5546 · a9a5546
2 parents 928f812 + a59ed84
commit a9a5546
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/finetune/train_cogvideox_lora.py b/finetune/train_cogvideox_lora.py
@@ -40,7 +40,7 @@
 from diffusers.pipelines.cogvideo.pipeline_cogvideox import get_resize_crop_region_for_grid
 from diffusers.training_utils import (
     cast_training_params,
-    clear_objs_and_retain_memory,
+    free_memory,
 )
 from diffusers.utils import check_min_version, convert_unet_state_dict_to_peft, export_to_video, is_wandb_available
 from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
@@ -726,7 +726,7 @@ def log_validation(
                 }
             )
 
-    clear_objs_and_retain_memory([pipe])
+    free_memory()
 
     return videos