Lightning 2.1.3 changes how checkpoints are saved, which breaks
how we handle saving and loading. This should be reverted in the
next release of lightning.

pytorch-lightning is in the process of changing their package names to
simply lightning.