From bc6051fdc2dad671ba5e7719febca9f2767442b3 Mon Sep 17 00:00:00 2001 From: Duoong Date: Sun, 8 Feb 2026 06:41:19 +0700 Subject: [PATCH] Fix W&B logging bugs --- anima_train.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/anima_train.py b/anima_train.py index 9fc1ec80..a86c30c3 100644 --- a/anima_train.py +++ b/anima_train.py @@ -578,6 +578,11 @@ def train(args): init_kwargs=init_kwargs, ) + if "wandb" in [tracker.name for tracker in accelerator.trackers]: + import wandb + wandb.define_metric("epoch") + wandb.define_metric("loss/epoch", step_metric="epoch") + if is_swapping_blocks: accelerator.unwrap_model(dit).prepare_block_swap_before_forward() @@ -782,8 +787,8 @@ def train(args): break if len(accelerator.trackers) > 0: - logs = {"loss/epoch": loss_recorder.moving_average} - accelerator.log(logs, step=epoch + 1) + logs = {"loss/epoch": loss_recorder.moving_average, "epoch": epoch + 1} + accelerator.log(logs, step=global_step) accelerator.wait_for_everyone()