From 2a188f07e682ed5dd958821a223d48c17a9aeb83 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 17 Nov 2024 16:12:10 +0900 Subject: [PATCH] Fix to work DOP with bock swap --- flux_train_network.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flux_train_network.py b/flux_train_network.py index 704c4d32e..679db62b6 100644 --- 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],