From 20e426cb5dcd4baae5f16fffc24cb10a81264577 Mon Sep 17 00:00:00 2001
From: Chengrui Wang <80876977+crywang@users.noreply.github.com>
Date: Sat, 22 Apr 2023 11:34:28 +0800
Subject: [PATCH] Fix bug in train_dreambooth_lora (#3183)

* Update train_dreambooth_lora.py

fix bug

* Update train_dreambooth_lora.py
---
 examples/dreambooth/train_dreambooth_lora.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py
index 1b75402c35..367a3422de 100644
--- a/examples/dreambooth/train_dreambooth_lora.py
+++ b/examples/dreambooth/train_dreambooth_lora.py
@@ -743,7 +743,7 @@ def main(args):
         )
         temp_pipeline._modify_text_encoder(text_lora_attn_procs)
         text_encoder = temp_pipeline.text_encoder
-        accelerator.register_for_checkpointing(unet_lora_layers)
+        accelerator.register_for_checkpointing(text_encoder_lora_layers)
         del temp_pipeline
 
     if args.scale_lr: