As a first step towards using lightning-fabric (or accelerator), we need to switch to using torch `DataLoader`s