Update

1 files changed, 6 insertions, 6 deletions

diff --git a/train_ti.py b/train_ti.py
index 775b918..f1dbed1 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -889,9 +889,6 @@ def main():
         seed=args.seed
     )
 
-    if accelerator.is_main_process:
-        checkpointer.save_samples(global_step_offset, args.sample_steps)
-
     local_progress_bar = tqdm(
         range(num_update_steps_per_epoch + num_val_steps_per_epoch),
         disable=not accelerator.is_local_main_process,
@@ -908,6 +905,10 @@ def main():
 
     try:
         for epoch in range(num_epochs):
+            if accelerator.is_main_process:
+                if epoch % args.sample_frequency == 0:
+                    checkpointer.save_samples(global_step + global_step_offset, args.sample_steps)
+
             local_progress_bar.set_description(f"Epoch {epoch + 1} / {num_epochs}")
             local_progress_bar.reset()
 
@@ -1001,13 +1002,11 @@ def main():
                         "global_step": global_step + global_step_offset
                     })
 
-                if (epoch + 1) % args.sample_frequency == 0:
-                    checkpointer.save_samples(global_step + global_step_offset, args.sample_steps)
-
         # 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.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
             })
@@ -1017,6 +1016,7 @@ 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
             })