Fix to work DOP with bock swap

2026-04-06 13:47:06 +00:00 · 2024-11-17 16:12:10 +09:00
parent 0047bb1fc3
commit 2a188f07e6
1 changed files with 1 additions and 0 deletions
--- a/flux_train_network.py
+++ b/flux_train_network.py
@@ -445,6 +445,7 @@ class FluxNetworkTrainer(train_network.NetworkTrainer):

            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],