fix: wait all block trasfer before siwtching offloader mode

2026-04-17 17:24:21 +00:00 · 2026-02-10 21:46:58 +09:00
parent 4992aae311
commit dbb40ae4c0
3 changed files with 10 additions and 0 deletions
--- a/library/anima_train_utils.py
+++ b/library/anima_train_utils.py
@@ -495,6 +495,7 @@ def sample_images(

    with torch.no_grad(), accelerator.autocast():
        for prompt_dict in prompts:
+            dit.prepare_block_swap_before_forward()
            _sample_image_inference(
                accelerator,
                args,