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.