diff --git a/train_network.py b/train_network.py index 7a589308..a661d31b 100644 --- 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