2 files changed, 2 insertions, 4 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()
diff --git a/train_ti.py b/train_ti.py
index f1dbed1..5cbef06 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -1006,7 +1006,7 @@ def main():
         if accelerator.is_main_process:
             print("Finished! Saving final checkpoint and resume state.")
             checkpointer.checkpoint(global_step + global_step_offset, "end")
-            checkpointer.save_samples(global_step_offset, args.sample_steps)
+            checkpointer.save_samples(global_step + global_step_offset, args.sample_steps)
             save_args(basepath, args, {
                 "global_step": global_step + global_step_offset
             })
@@ -1016,7 +1016,6 @@ def main():
         if accelerator.is_main_process:
             print("Interrupted, saving checkpoint and resume state...")
             checkpointer.checkpoint(global_step + global_step_offset, "end")
-            checkpointer.save_samples(global_step_offset, args.sample_steps)
             save_args(basepath, args, {
                 "global_step": global_step + global_step_offset
             })