Unverified Commit 677df5ac authored by Shyam Marjit's avatar Shyam Marjit Committed by GitHub
Browse files

fixed SDXL text encoder training bug #5016 (#5078)


Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
parent 16851efa
...@@ -1070,6 +1070,11 @@ def main(args): ...@@ -1070,6 +1070,11 @@ def main(args):
if args.train_text_encoder: if args.train_text_encoder:
text_encoder_one.train() text_encoder_one.train()
text_encoder_two.train() text_encoder_two.train()
# set top parameter requires_grad = True for gradient checkpointing works
text_encoder_one.text_model.embeddings.requires_grad_(True)
text_encoder_two.text_model.embeddings.requires_grad_(True)
for step, batch in enumerate(train_dataloader): for step, batch in enumerate(train_dataloader):
with accelerator.accumulate(unet): with accelerator.accumulate(unet):
pixel_values = batch["pixel_values"].to(dtype=vae.dtype) pixel_values = batch["pixel_values"].to(dtype=vae.dtype)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment