From 8ba02ac8296c08ef8c880d87d7134504550671cb Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sat, 22 Jul 2023 09:56:36 +0900 Subject: [PATCH] fix to work text encoder only network with bf16 --- train_network.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/train_network.py b/train_network.py index b482c80a..6f41d199 100644 --- 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