Write the model in a parallel framework which we do not need to maintain

I found this library: 
https://github.com/williamFalcon/pytorch-lightning

As I understand it, they make it easy to parallelize torch models.