From aea7cada42cbb3373fc04bf98f7cae0a91d77d27 Mon Sep 17 00:00:00 2001
From: DKnight54 <126916963+DKnight54@users.noreply.github.com>
Date: Sun, 2 Feb 2025 04:54:36 +0800
Subject: [PATCH] Update train_network.py

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

diff --git a/train_network.py b/train_network.py
index 5d170926..17a3b36e 100644
--- a/train_network.py
+++ b/train_network.py
@@ -1024,7 +1024,7 @@ class NetworkTrainer:
 
                 # Checks if the accelerator has performed an optimization step behind the scenes
                 # Collecting latents and caption lists from all processes
-                all_latents = gather_object(latents)
+                all_latents = gather_object([latents])
                 all_captions =  gather_object(batch["captions"])
                 logger.info(f"latents: {latents}")
                 logger.info(f"all_latents: {all_latents}")