call optimizer eval/train for sample_at_first, also set train after resuming closes #1667

2026-04-09 06:45:09 +00:00 · 2024-10-04 20:35:16 +09:00
parent 8bea039a8d
commit ba08a89894
2 changed files with 4 additions and 0 deletions
--- a/flux_train.py
+++ b/flux_train.py
@@ -706,7 +706,9 @@ def train(args):
        accelerator.unwrap_model(flux).prepare_block_swap_before_forward()

    # For --sample_at_first
+    optimizer_eval_fn()
    flux_train_utils.sample_images(accelerator, args, 0, global_step, flux, ae, [clip_l, t5xxl], sample_prompts_te_outputs)
+    optimizer_train_fn()
    if len(accelerator.trackers) > 0:
        # log empty object to commit the sample images to wandb
        accelerator.log({}, step=0)