[DreamBooth] Set train mode for text encoder (#1012)

Set train mode for text encoder

[DreamBooth] Set train mode for text encoder (#1012)
Set train mode for text encoder
4623f095 · Duong A. Nguyen · GitHub · abe05822 · 4623f095
Unverified Commit 4623f095 authored Oct 27, 2022 by Duong A. Nguyen Committed by GitHub Oct 27, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

examples/dreambooth/train_dreambooth.py examples/dreambooth/train_dreambooth.py +2 -0

No files found.
--- a/examples/dreambooth/train_dreambooth.py
+++ b/examples/dreambooth/train_dreambooth.py
@@ -574,6 +574,8 @@ def main(args):

    for epoch in range(args.num_train_epochs):
        unet.train()
+        if args.train_text_encoder:
+            text_encoder.train()
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(unet):
                # Convert images to latent space