scvi-tools
scvi-tools copied to clipboard
Add support for DDP.
This PR is a draft and not completed, but it shows a slightly different approach which is compatible with DDP.
Explanation: PyTorch Lightning injects a DistributedSampler within the DataLoader. To do so, it tries to access init arguments from dataloader attribues.
With the current custom AnnDataLoader, AnnDataset isn't saved, and therefore Lightning will fail.
To avoid further bug, I would recommend to avoid overriding the DataLoader unless there is a strong requirement or at least make sure it can re-instantianted from its attributes.
Best regards, T.C