1 files changed, 6 insertions, 8 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 48bdcf8..9c1e41c 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -1,6 +1,7 @@
 import argparse
 import datetime
 import logging
+import itertools
 from pathlib import Path
 from functools import partial
 
@@ -578,14 +579,11 @@ def main():
     datamodule.setup()
 
     optimizer = optimizer_class(
-        [
-            {
-                'params': unet.parameters(),
-            },
-            {
-                'params': text_encoder.parameters(),
-            }
-        ],
+        itertools.chain(
+            unet.parameters(),
+            text_encoder.text_model.encoder.parameters(),
+            text_encoder.text_model.final_layer_norm.parameters(),
+        ),
         lr=args.learning_rate,
         betas=(args.adam_beta1, args.adam_beta2),
         weight_decay=args.adam_weight_decay,