Merge pull request #2338 from haofanwang/patch-1

Fix a typo in train_dreambooth_colossalai.py

Merge pull request #2338 from haofanwang/patch-1
Fix a typo in train_dreambooth_colossalai.py
7a332b17 · Fazzie-Maqianli · GitHub · 8b1e0dfd · 9edd0aa7 · 7a332b17
Unverified Commit 7a332b17 authored Jan 06, 2023 by Fazzie-Maqianli Committed by GitHub Jan 06, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/images/dreambooth/train_dreambooth_colossalai.py examples/images/dreambooth/train_dreambooth_colossalai.py +1 -1

No files found.
--- a/examples/images/dreambooth/train_dreambooth_colossalai.py
+++ b/examples/images/dreambooth/train_dreambooth_colossalai.py
@@ -484,7 +484,7 @@ def main(args):
        unet.enable_gradient_checkpointing()

    if args.scale_lr:
-        args.learning_rate = args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * 2
+        args.learning_rate = args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * gpc.get_world_size(ParallelMode.DATA)

    unet = gemini_zero_dpp(unet, pg, args.placement)