github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 18:19:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

modify shell for check

Browse Source
This commit is contained in:
Maruyama_Aya
2023-06-08 13:17:58 +08:00
parent cf4792c975
commit 039854b391
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -487,7 +487,7 @@ def main(args):
if args.plugin.startswith('torch_ddp'):
plugin = TorchDDPPlugin()
elif args.plugin == 'gemini':
plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2 ** 5)
plugin = GeminiPlugin(placement_policy=args.placement, strict_ddp_mode=True, initial_scale=2 ** 5)
elif args.plugin == 'low_level_zero':
plugin = LowLevelZeroPlugin(initial_scale=2 ** 5)