Integrate with LightningLite

[Aniket Maurya]

Hi @junyanz, we love this project! Easiest way to get started with GANs 🔥

I noticed that you're using a basic PyTorch Training loop. With a small refactor it can be enhanced for better support on any kind of device as well as distributed training while you retain full control over loops and optimization logic. You can check out our blog post about it here: Scale your PyTorch code with LightningLite.

Feel free to get in touch and let us know what you think about this, and if you need any help with the PR. Thanks! 😄