Fix keys for lora format on advanced training scripts (huggingface#6361)

fix keys for lora format on advanced training scripts
donhardman · Dec 29, 2023 · f3b1509 · f3b1509
1 parent e192db3
commit f3b1509
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py
@@ -1955,7 +1955,7 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers):
     if accelerator.is_main_process:
         unet = accelerator.unwrap_model(unet)
         unet = unet.to(torch.float32)
-        unet_lora_layers = get_peft_model_state_dict(unet)
+        unet_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
 
         if args.train_text_encoder:
             text_encoder_one = accelerator.unwrap_model(text_encoder_one)