Replace TensroFlow with PyTorch

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)