For a custom dataset, you can inherit torch.utils.data.Dataset class with
len(returns length of your dataset), getitem(takes in the index) and you would need to use some files libraries and image ones to open and load files from your custom dataset.
You can check these links out:

https://pytorch.org/tutorials/beginner/data_loading_tutorial.html 3
https://stanford.edu/~shervine/blog/pytorch-how-to-generate-data-parallel 1

Hopefully, these links will be helpful.