Draft: Training plots/graphs + rudimentary model checkpointing after every epoch

Pull Request Title

Add the options to save matplotlib plots of learning rate and train/validation losses during training.

Description

Wants to merge: feature/plots into main

Type of change

• Bug fix
• New feature
• Enhancement
• Documentation update
• Other (specify right below)

Merge request commits

• Save loss data points, refactor into functions

Refactor the plots and model saving code into separate functions.

Save the loss and learning rate scheduling data points to .pth files, in addition to saving the graphs. This is helpful to compare different training/valid loss curves of different models post-training, by creating custom graphs to represent all the information together.

• Add train plots, add rudimentary training checkpointing

Add the option to save training data points such as learning rate and train/valid losses as graphs/plots.

Add a very rudimentary training checkpointing mechanism, which currently only gives the user the option to save the model every epoch.

Related Issues

Screenshots or GIFs

Checklist

• I have tested the code with the changes manually.
• My code follows the project's style guidelines.
• I have documented my code for others to understand.
• I have updated documentation as needed (including README.md, code comments and doc strings).

Reviewer Guidelines

Additional Notes

@mentions