Saving checkpoints during the model training (PyTorch backend)
Problem: The PyTorch backend allows saving the trained model only after the complete training.
Improvement: The training pipeline should save the model checkpoints after each epoch. If one tries to run the incomplete experiment later, there should be a provision in the training pipeline to load the latest checkpoint and continue the training from there.