github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update train_dreambooth_colossalai.py 

accelerator.num_processes -> gpc.get_world_size(ParallelMode.DATA)
pull/2338/head
Haofan Wang 2023-01-05 15:49:57 +08:00 committed by GitHub
parent f1bc2418c4
commit 9edd0aa75e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
examples/images/dreambooth/train_dreambooth_colossalai.py
View File

@ -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)