Search code examples
inheritancepytorchoverridingpytorch-dataloader

Inheriting the PyTorch DataLoader class

Can we inherit the DataLoader class? If so, are there any specific restrictions to it?
I know we can do so for the Dataset class, but I need to know specifically about the DataLoader.

I have some requirements for a specific use case which involves overriding some of the DataLoader methods and adding some new ones.

Would really appreciate the help! Thanks....

Solution

  • Thanks @Jan and @TheodorPeifer for your useful inputs..! :)

    I was able to resolve this with a combination of NamedTuple returns from the dataset __get_item__ calls with the "worker_specific_info"(torch.utils.data.get_worker_info()) along with a custom collate_fn implementation.