[example] fixed seed error in train_dreambooth_colossalai.py (#2445)

cfd1d5ee · Haofan Wang · GitHub · ac18a445 · cfd1d5ee
Unverified Commit cfd1d5ee authored Jan 11, 2023 by Haofan Wang Committed by GitHub Jan 11, 2023
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 3 deletions

examples/images/dreambooth/train_dreambooth_colossalai.py examples/images/dreambooth/train_dreambooth_colossalai.py +4 -3

No files found.
--- a/examples/images/dreambooth/train_dreambooth_colossalai.py
+++ b/examples/images/dreambooth/train_dreambooth_colossalai.py
@@ -355,10 +355,11 @@ def gemini_zero_dpp(model: torch.nn.Module, placememt_policy: str = "auto"):


 def main(args):
-    colossalai.launch_from_torch(config={})

-    if args.seed is not None:
-        gpc.set_seed(args.seed)
+    if args.seed is None:
+        colossalai.launch_from_torch(config={})
+    else:
+        colossalai.launch_from_torch(config={}, seed=args.seed)

    if args.with_prior_preservation:
        class_images_dir = Path(args.class_data_dir)