Would be nice to have example with PyTorch Lightning
Would be nice to have example with PyTorch Lightning