From b3b856393cdebc355ab8e631553b09736d5bfe30 Mon Sep 17 00:00:00 2001 From: DKnight54 <126916963+DKnight54@users.noreply.github.com> Date: Mon, 3 Feb 2025 03:38:24 +0800 Subject: [PATCH] Update train_network.py --- train_network.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/train_network.py b/train_network.py index 2accc40a..9364cc19 100644 --- a/train_network.py +++ b/train_network.py @@ -39,7 +39,7 @@ from library.custom_train_functions import ( apply_masked_loss, ) from library.utils import setup_logging, add_logging_arguments -from accelerate.utils import gather_object +from accelerate.utils import gather_object, gather setup_logging() import logging @@ -1024,10 +1024,11 @@ 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]) + logger.info(f"latents.size: {latents.size()} before gather on device {accelerator.state.local_process_index}") + all_latents = gather(latents) all_captions = gather_object(batch["captions"]) #logger.info(f"latents: {latents}") - #logger.info(f"all_latents: {all_latents}") + logger.info(f"all_latents.size: {all_latents.size()}") example_tuple = (all_latents, all_captions) if accelerator.sync_gradients: progress_bar.update(1)