You can define your custom dataset class extending the torch.utils.data.Dataset class. You have to implement the init(), getitem() and len() classes. Here is a good tutorial for writing custom datasets Pytorch Custom Dataset 3