Since your batch size is set to 100 in the DataLoader, you shouldn’t get 10000 or 10001 samples in a single batch.
Could you post the complete code by wrapping it into three backticks ```?
Also, since you are already dealing with images, you don’t need to flatten it inside the CustomDataset and the n reshaping it back to the image tensor shape.