Update to work with PyTorch lightning

[Jumperkables]

Consider this, we can better trust train/val code and metrics with no bugs. Would make the different things very scalable, and even easier to pop a transformer in there