Compatibility issue with PyTorch Lightning >= 1.6.0

[rasbt]

Hey, I was just thinking about implementing k-fold support and saw this awesome library! Great work! Just wanted to drop you a note that I noticed that it currently doesn't work with PyTorch Lightning >= 1.6.0. It works fine with e.g., 1.5.10 but in 1.6.0 it throws an error when calling trainer.cross_validate():

It happens in both cases whether you use a datamodule or train_dataloader.