Fix scheduler warning: use last_epoch instead of step() loop on resume

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 18:23:17 +03:00
parent 905b9867c8
commit 517b87d3d8

View File

@@ -331,9 +331,8 @@ def train(cfg: TrainConfigGTAUAV) -> None:
if "loss_state" in resume_ckpt:
loss_fn.load_state_dict(resume_ckpt["loss_state"])
LOGGER.info("🔄 Loss state restored (tau=%.4f)", loss_fn.current_temperature)
# Advance scheduler to the correct step.
for _ in range(start_epoch * steps_per_epoch):
scheduler.step()
# Set scheduler last_epoch so it resumes at the correct LR.
scheduler.last_epoch = start_epoch * steps_per_epoch
LOGGER.info("🔄 Resuming from epoch %d", start_epoch)
history: list[dict] = []