How can I resume the state of a stateful dataloader in pytorch lightning ?

[hiyyg]

I want to use the stateful dataloader in pytorch lightning.
https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader

What should I do to make it work?

[nrxszvo]

Just drop-in replace the standard torch DataLoader in your LightningDataModule with the StatefulDataLoader from torchdata:

def train_dataloader(self):
        return StatefulDataLoader(
            self.trainset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
        )

Note that if you use a progress bar, you might need to set the total number of training batches explicitly (I get a ? if I don't), e.g.:

class MyProgBar(TQDMProgressBar):
    def init_train_tqdm(self):
        bar = super().init_train_tqdm()
        bar.total = self.trainer.num_training_batches
        return bar