Fixes LoRA SDXL training script with DDP + PEFT (#6816)

Update train_dreambooth_lora_sdxl.py
huggingface · Feb 3, 2024 · 15ed53d · 15ed53d
1 parent 9cc59ba
commit 15ed53d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py
@@ -1399,8 +1399,8 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers):
             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)
+            accelerator.unwrap_model(text_encoder_one).text_model.embeddings.requires_grad_(True)
+            accelerator.unwrap_model(text_encoder_two).text_model.embeddings.requires_grad_(True)
 
         for step, batch in enumerate(train_dataloader):
             with accelerator.accumulate(unet):