Port the Tensorflow training code to a more modern PyTorch Lightning-based workflow. Current training class - [link](https://github.com/iris-hep/NSBI-workflow-tutorial/blob/main/src/nsbi_common_utils/training.py)