You should create a custom dataset class, which loads the corresponding csv file and return the tuple <image, label> in the __getitem__ function: Writing Custom Datasets, DataLoaders and Transforms — PyTorch Tutorials 1.8.1+cu102 documentation 3