From 18a6059d749e2974f7313eac9439f562a8176f59 Mon Sep 17 00:00:00 2001
From: sdbds <865105819@qq.com>
Date: Wed, 5 Jun 2024 20:12:36 +0800
Subject: [PATCH] Add DDP transform

---
 sdxl_train_control_net_lllite.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/sdxl_train_control_net_lllite.py b/sdxl_train_control_net_lllite.py
index 301310901..2f8b853fa 100644
--- a/sdxl_train_control_net_lllite.py
+++ b/sdxl_train_control_net_lllite.py
@@ -289,6 +289,8 @@ def train(args):
     # acceleratorがなんかよろしくやってくれるらしい
     unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
 
+    unet = unet.module if isinstance(unet, DDP) else unet
+
     if args.gradient_checkpointing:
         unet.train()  # according to TI example in Diffusers, train is required -> これオリジナルのU-Netしたので本当は外せる
     else: