fix: save checkpoints after optimizer steps

- add a post-step callback hook for checkpoint saves - preserve updated model, optimizer, and scheduler state - cover checkpoint ordering with a regression test
2026-09-01 12:25:27 +08:00
parent 432dfec3c2
commit 08721f6d31
3 changed files with 42 additions and 2 deletions
@@ -101,6 +101,8 @@ class Trainer:
                            if context.scheduler:
                                context.scheduler.step()

+                            self._call_callbacks("on_after_optimizer_step", context)
+
                self._call_callbacks("on_epoch_end", context)

            if context.stop_requested: