fix Text Encoder only LoRA training

kohya-ss · Jan 27, 2025 · 0778dd9 · 0778dd9
1 parent 59b3b94
commit 0778dd9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/flux_train_network.py b/flux_train_network.py
@@ -378,7 +378,7 @@ def get_noise_pred_and_target(
         def call_dit(img, img_ids, t5_out, txt_ids, l_pooled, timesteps, guidance_vec, t5_attn_mask):
             # if not args.split_mode:
             # normal forward
-            with torch.set_grad_enabled(is_train and train_unet), accelerator.autocast():
+            with torch.set_grad_enabled(is_train), accelerator.autocast():
                 # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transformer model (we should not keep it but I want to keep the inputs same for the model for testing)
                 model_pred = unet(
                     img=img,

diff --git a/sd3_train_network.py b/sd3_train_network.py
@@ -345,7 +345,7 @@ def get_noise_pred_and_target(
             t5_attn_mask = None
 
         # call model
-        with torch.set_grad_enabled(is_train and train_unet), accelerator.autocast():
+        with torch.set_grad_enabled(is_train), accelerator.autocast():
             # TODO support attention mask
             model_pred = unet(noisy_model_input, timesteps, context=context, y=lg_pooled)
 

diff --git a/train_network.py b/train_network.py
@@ -233,7 +233,7 @@ def get_noise_pred_and_target(
                 t.requires_grad_(True)
 
         # Predict the noise residual
-        with torch.set_grad_enabled(is_train and train_unet), accelerator.autocast():
+        with torch.set_grad_enabled(is_train), accelerator.autocast():
             noise_pred = self.call_unet(
                 args,
                 accelerator,