Update

1 files changed, 7 insertions, 4 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index f8f6e84..a85ae4c 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -582,12 +582,15 @@ def main():
     )
     datamodule.setup()
 
-    optimizer = create_optimizer(
-        itertools.chain(
-            unet.parameters(),
+    params_to_optimize = (unet.parameters(), )
+    if args.train_text_encoder_epochs != 0:
+        params_to_optimize += (
             text_encoder.text_model.encoder.parameters(),
             text_encoder.text_model.final_layer_norm.parameters(),
-        ),
+        )
+
+    optimizer = create_optimizer(
+        itertools.chain(*params_to_optimize),
         lr=args.learning_rate,
     )