Fix

1 files changed, 1 insertions, 2 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index c0fe328..05f6cb5 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -1068,14 +1068,13 @@ def main():
         # Create the pipeline using using the trained modules and save it.
         if accelerator.is_main_process:
             print("Finished! Saving final checkpoint and resume state.")
-            checkpointer.save_samples(0, args.sample_steps)
+            checkpointer.save_samples(global_step, args.sample_steps)
             checkpointer.save_model()
             accelerator.end_training()
 
     except KeyboardInterrupt:
         if accelerator.is_main_process:
             print("Interrupted, saving checkpoint and resume state...")
-            checkpointer.save_samples(0, args.sample_steps)
             checkpointer.save_model()
             accelerator.end_training()
         quit()