Fix issues for saving checkpointing steps

leisuzz · Nov 13, 2024 · 486a115 · 486a115
1 parent 6837a81
commit 486a115
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/examples/dreambooth/train_dreambooth_flux.py b/examples/dreambooth/train_dreambooth_flux.py
@@ -1187,7 +1187,8 @@ def save_model_hook(models, weights, output_dir):
                     raise ValueError(f"Wrong model supplied: {type(model)=}.")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
     def load_model_hook(models, input_dir):
         for _ in range(len(models)):

diff --git a/examples/dreambooth/train_dreambooth_lora_flux.py b/examples/dreambooth/train_dreambooth_lora_flux.py
@@ -1250,7 +1250,8 @@ def save_model_hook(models, weights, output_dir):
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
             FluxPipeline.save_lora_weights(
                 output_dir,