Implement training on stages

[Erlemar]

Often it is a good idea to train model on stages - for example with changes of loss, augs or other things. Need to implement it and be able to set it using configs.

https://forums.pytorchlightning.ai/t/what-is-the-best-way-to-train-on-stages/95/4