fix to work text encoder only network with bf16

Kuron88 · Jul 22, 2023 · 8ba02ac · 8ba02ac
1 parent 73a08c0
commit 8ba02ac
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/train_network.py b/train_network.py
@@ -401,6 +401,8 @@ def train(self, args):
                     text_encoder, network, optimizer, train_dataloader, lr_scheduler
                 )
                 text_encoders = [text_encoder]
+
+            unet.to(accelerator.device, dtype=weight_dtype) # move to device because unet is not prepared by accelerator
         else:
             network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
                 network, optimizer, train_dataloader, lr_scheduler