Skip the final 1 step

feffy380 · May 31, 2024 · 3eb27ce · 3eb27ce
1 parent b2363f1
commit 3eb27ce
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/train_network.py b/train_network.py
@@ -897,6 +897,10 @@ def remove_model(old_ckpt_name):
 
             for step, batch in enumerate(skipped_dataloader or train_dataloader):
                 current_step.value = global_step
+                if initial_step > 0:
+                    initial_step -= 1
+                    continue
+
                 with accelerator.accumulate(training_model):
                     on_step_start(text_encoder, unet)