From 2cdaa33147e75e7f7dd662f060eed37690321ae6 Mon Sep 17 00:00:00 2001 From: DKnight54 <126916963+DKnight54@users.noreply.github.com> Date: Fri, 31 Jan 2025 21:47:15 +0800 Subject: [PATCH] Update train_network.py --- train_network.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/train_network.py b/train_network.py index 5f2a31da..7563424b 100644 --- a/train_network.py +++ b/train_network.py @@ -1023,14 +1023,8 @@ class NetworkTrainer: # Checks if the accelerator has performed an optimization step behind the scenes # Collecting latents and caption lists from all processes - all_lists_of_latents = gather_object(latents) - all_lists_of_captions = gather_object(batch["captions"]) - all_latents = [] - all_captions = [] - for ilatents in all_lists_of_latents: - all_latents.extend(ilatents) - for icaptions in all_lists_of_captions: - all_captions.extend(icaptions) + all_latents = gather_object(latents) + all_captions = gather_object(batch["captions"]) example_tuple = (all_latents, all_captions) if accelerator.sync_gradients: progress_bar.update(1)