From bf2e5abd2eb5b395c1cc2cc142bf64e8e0e764d6 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 3 Apr 2025 17:57:28 -0400 Subject: [PATCH] Move batch size syncing outside validation --- train_network.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train_network.py b/train_network.py index 55be9601..885203cf 100644 --- a/train_network.py +++ b/train_network.py @@ -1566,9 +1566,6 @@ class NetworkTrainer: } self.step_logging(accelerator, logs, global_step, epoch=epoch + 1) - if accelerator.sync_gradients: - batch_size = 0 # reset batch size - restore_rng_state(rng_states) args.min_timestep = original_args_min_timestep args.max_timestep = original_args_max_timestep @@ -1576,6 +1573,9 @@ class NetworkTrainer: accelerator.unwrap_model(network).train() progress_bar.unpause() + if accelerator.sync_gradients: + batch_size = 0 # reset batch size + if global_step >= args.max_train_steps: break