Don't rely on Accelerate for gradient accumulation

1 files changed, 1 insertions, 1 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 431ff3d..280cf77 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -439,7 +439,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
     )
 
@@ -590,6 +589,7 @@ def main():
         lr_scheduler=lr_scheduler,
         prepare_unet=True,
         num_train_epochs=args.num_train_epochs,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
         sample_frequency=args.sample_frequency,
         # --
         tokenizer=tokenizer,