Add optional checkpoint saving to Trainer class

[blkdmr]

The Trainer class in fenn/nn/trainers currently provides a clean training loop abstraction, but it does not offer a built-in way to save model checkpoints during training.

Checkpointing is a common requirement for:

• long training runs
• recovering from interruptions
• inspecting intermediate models
• selecting the best-performing epoch

This issue proposes adding a simple, optional checkpoint saving mechanism to the Trainer, without changing default behavior.

Goal

Allow users to optionally enable checkpoint saving by specifying:

• which epochs should be saved
• a base name for the checkpoint files

If not enabled, the Trainer should behave exactly as it does now.

Proposed behavior

Extend the Trainer configuration with optional parameters such as:

• checkpoint_epochs: int | None
  • epochs at which a checkpoint is saved (e.g. 5, 10, 20)
• checkpoint_name: str | None
  • base filename used for checkpoints

When enabled:

• at the end of a specified epoch, the trainer saves:
  • model state (state_dict)
  • optimizer state (if available)
  • current epoch index

Example filenames:

• checkpoint_name_epoch_10.pt
• checkpoint_name_epoch_20.pt

Tasks

1. Inspect the current Trainer implementation in fenn/nn/trainers.
2. Add checkpoint-related arguments (constructor or config.yaml-based).
3. Implement checkpoint saving at the end of training epochs.
4. Ensure:
   • no checkpoints are written unless explicitly enabled
   • existing training workflows remain unaffected
5. Add minimal documentation or docstring explaining usage.

Acceptance criteria

• Checkpoint saving is fully optional.
• Users can control:
  • which epochs are saved
  • the checkpoint file base name
• Trainer behavior is unchanged when checkpointing is disabled.
• Saved checkpoints can be reloaded with standard PyTorch APIs.
• Code remains simple and readable.

Optional (but nice to have)

• resume-from-checkpoint logic

How to contribute

Comment on this issue to claim it, then open a PR with:

• the implementation
• a short usage example (code snippet or docstring)

[spencerirving]

I'd like to work on this issue if it is still available!

[blkdmr]

Sure @spencerirving! If you want, you can also add the possibility to pass a validation set to the trainer. In this case, the user can dedice to save the "latest" checkpoint or the "best", based on the performances on the validation set.

[spencerirving]

Sorry this has been taking a while school and career fairs just kept me super busy. I plan to work on it this weekend.

[blkdmr]

Don’t worry about it, and thanks for letting me know!

[blkdmr]

Hi @spencerirving, how’s it going? Do you have any updates? Feel free to ask for help if you need any.

[spencerirving]

Hey @blkdmr,

Sorry for taking so long with this, a new semester of school just started for me and I didn't realize how much work I had and unfortunately this project got lost in the mist. I just opened a pull request with the implementation. Please let me know what you think and feel free to suggest any additional changes and features!

I had fun working with this project and hope to find more time in the future to help with future issues!