ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: set tqdm postfix before update so first step shows metrics

This commit is contained in:
ViperEkura 2026-07-07 00:14:13 +08:00
parent 4af7acd449
commit 536dbc0c9a
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
astrai/trainer/train_callback.py
View File

@ -209,7 +209,6 @@ class ProgressBarCallback(TrainCallback):
@only_on_rank(0) @only_on_rank(0)
def on_optimizer_step(self, context: TrainContext): def on_optimizer_step(self, context: TrainContext):
self.progress_bar.update(1)
postfix = { postfix = {
"step": context.optimizer_step, "step": context.optimizer_step,
"loss": f"{context.loss:.4f}", "loss": f"{context.loss:.4f}",
@ -220,6 +219,7 @@ class ProgressBarCallback(TrainCallback):
if context.val_loss is not None: if context.val_loss is not None:
postfix["val_loss"] = f"{context.val_loss:.4f}" postfix["val_loss"] = f"{context.val_loss:.4f}"
self.progress_bar.set_postfix(postfix) self.progress_bar.set_postfix(postfix)
self.progress_bar.update(1)
@only_on_rank(0) @only_on_rank(0)
def on_epoch_end(self, context: TrainContext): def on_epoch_end(self, context: TrainContext):