Fix to work DOP with bock swap

kohya-ss · Nov 17, 2024 · 2a188f0 · 2a188f0
1 parent 0047bb1
commit 2a188f0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/flux_train_network.py b/flux_train_network.py
@@ -445,6 +445,7 @@ def call_dit(img, img_ids, t5_out, txt_ids, l_pooled, timesteps, guidance_vec, t
 
             if len(diff_output_pr_indices) > 0:
                 network.set_multiplier(0.0)
+                unet.prepare_block_swap_before_forward()
                 with torch.no_grad():
                     model_pred_prior = call_dit(
                         img=packed_noisy_model_input[diff_output_pr_indices],