set static graph flag when DDP ref kohya-ss#1363

feffy380 · Jun 9, 2024 · 58fb648 · 58fb648
1 parent e5bab69
commit 58fb648
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/sdxl_train_control_net_lllite.py b/sdxl_train_control_net_lllite.py
@@ -289,6 +289,9 @@ def train(args):
     # acceleratorがなんかよろしくやってくれるらしい
     unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
 
+    if isinstance(unet, DDP):
+        unet._set_static_graph() # avoid error for multiple use of the parameter
+
     if args.gradient_checkpointing:
         unet.train()  # according to TI example in Diffusers, train is required -> これオリジナルのU-Netしたので本当は外せる
     else: