Fix bug in train_dreambooth_lora (#3183)

crywang · web-flow · commit 20e426cb5dcd · 2023-04-22T09:04:28.000+05:30
* Update train_dreambooth_lora.py

fix bug

* Update train_dreambooth_lora.py
diff --git 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:

Original file line number	Diff line number	Diff line change
`@@ -743,7 +743,7 @@ def main(args):`
`743`	`743`	`)`
`744`	`744`	`temp_pipeline._modify_text_encoder(text_lora_attn_procs)`
`745`	`745`	`text_encoder = temp_pipeline.text_encoder`
`746`		`- accelerator.register_for_checkpointing(unet_lora_layers)`
	`746`	`+ accelerator.register_for_checkpointing(text_encoder_lora_layers)`
`747`	`747`	`del temp_pipeline`
`748`	`748`
`749`	`749`	`if args.scale_lr:`