Update train_network.py

2026-04-16 00:49:40 +00:00 · 2025-05-22 18:21:19 +08:00
parent 93ce543434
commit 115b1d0aa7
1 changed files with 1 additions and 0 deletions
--- a/train_network.py
+++ b/train_network.py
@@ -1096,6 +1096,7 @@ class NetworkTrainer:
                        train_util.save_and_remove_state_on_epoch_end(args, accelerator, epoch + 1)

            if args.sample_every_n_epochs is not None and (epoch + 1)% args.sample_every_n_epochs == 0:
+                accelerator.wait_for_everyone()
                self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, example_tuple)

            # end of epoch