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

Set train mode for text encoder
aravind-h-v · Oct 27, 2022 · 4623f09 · 4623f09
1 parent abe0582
commit 4623f09
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
@@ -574,6 +574,8 @@ def collate_fn(examples):
 
  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