fix placing of requires_grad_ of U-Net

kohya-ss · Oct 1, 2023 · 4cc9196 · 4cc9196
1 parent 81419f7
commit 4cc9196
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/train_network.py b/train_network.py
@@ -12,10 +12,13 @@
 
 from tqdm import tqdm
 import torch
+
 try:
     import intel_extension_for_pytorch as ipex
+
     if torch.xpu.is_available():
         from library.ipex import ipex_init
+
         ipex_init()
 except Exception:
     pass
@@ -428,8 +431,10 @@ def train(self, args):
                 # set top parameter requires_grad = True for gradient checkpointing works
                 if train_text_encoder:
                     t_enc.text_model.embeddings.requires_grad_(True)
-                else:
-                    unet.parameters().__next__().requires_grad_(True)
+
+            # set top parameter requires_grad = True for gradient checkpointing works
+            if not train_text_encoder:  # train U-Net only
+                unet.parameters().__next__().requires_grad_(True)
         else:
             unet.eval()
             for t_enc in text_encoders: