Fix to work training U-Net only LoRA for SD1/2

kohya-ss · Oct 1, 2023 · 81419f7 · 81419f7
1 parent 6bd6cd9
commit 81419f7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/train_network.py b/train_network.py
@@ -426,7 +426,10 @@ def train(self, args):
                 t_enc.train()
 
                 # set top parameter requires_grad = True for gradient checkpointing works
-                t_enc.text_model.embeddings.requires_grad_(True)
+                if train_text_encoder:
+                    t_enc.text_model.embeddings.requires_grad_(True)
+                else:
+                    unet.parameters().__next__().requires_grad_(True)
         else:
             unet.eval()
             for t_enc in text_encoders: