VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix to work text encoder only network with bf16

Browse Source
This commit is contained in:
Kohya S
2023-07-22 09:56:36 +09:00
parent 73a08c0be0
commit 8ba02ac829
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
train_network.py
View File

@@ -401,6 +401,8 @@ class NetworkTrainer:
text_encoder, network, optimizer, train_dataloader, lr_scheduler text_encoder, network, optimizer, train_dataloader, lr_scheduler
) )
text_encoders = [text_encoder] text_encoders = [text_encoder]
unet.to(accelerator.device, dtype=weight_dtype) # move to device because unet is not prepared by accelerator
else: else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler network, optimizer, train_dataloader, lr_scheduler