Save checkpoint on train epoch end

[mozharovsky]

This PR adds a new feature that saves checkpoints when training epoch is over.

Patch Notes

• Added a feature to save the checkpoints on training epoch end
• Updated Lightning to v1.0
• Minor tensor fixes with metrics logging