fix to work text encoder only network with bf16

2026-04-08 22:35:09 +00:00 · 2023-07-22 09:56:36 +09:00
parent 73a08c0be0
commit 8ba02ac829
1 changed files with 2 additions and 0 deletions
--- a/train_network.py
+++ b/train_network.py
@@ -401,6 +401,8 @@ class NetworkTrainer:
                    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