Don't rely on Accelerate for gradient accumulation

1 files changed, 1 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index 6dc07dd..68783ea 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -518,7 +518,6 @@ def main():
     accelerator = Accelerator(
         log_with=LoggerType.TENSORBOARD,
         logging_dir=f"{output_dir}",
-        gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision
     )
 
@@ -611,6 +610,7 @@ def main():
         low_freq_noise=0,
         strategy=textual_inversion_strategy,
         num_train_epochs=args.num_train_epochs,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
         sample_frequency=args.sample_frequency,
         checkpoint_frequency=args.checkpoint_frequency,
         milestone_checkpoints=not args.no_milestone_checkpoints,