Thoughts about integration with pytorch-lightning?

Great work here havakv, looks really modular and well thought out.

I'm interested in using some of the time-to-event tools here with high dimensional imaging inputs. I've been building up a medical imaging codebase using [pytorch-lightning](https://github.com/PyTorchLightning/pytorch-lightning) for a little bit mostly because of how modular & convenient it makes iterating over multiple experiments on a cluster environment.

Do you have any ideas of how best to re-organize some of pycox.models (PCHazard for example) to torch-lightning before I start on this? What I'm really interested in is being able to use the pytorch lightning [trainer](https://pytorch-lightning.readthedocs.io/en/stable/trainer.html).

Start by reviewing the pycox.models organization and the PCHazard example, then read the linked PyTorch Lightning Trainer documentation. The issue needs an agreed design for reorganizing the models around the Trainer; completion is not defined beyond establishing a workable integration approach.

Written by the indexing model from the issue text.