I noticed that ImageDataset could not be imported. Thats why, I added "from datasets import ImageDataset".