Add a util function for training loop

[SarahAlidoost]

Right now, training loop is defined inside the notebook. We want a util function for that:

• with a proper logging instead of printing of results.
• saving the model config and model itself