From 115b1d0aa753b54f4a74352f445d35d5b8eccb8b Mon Sep 17 00:00:00 2001
From: DKnight54 <126916963+DKnight54@users.noreply.github.com>
Date: Thu, 22 May 2025 18:21:19 +0800
Subject: [PATCH] Update train_network.py

---
 train_network.py | 1 +
 1 file changed, 1 insertion(+)

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