Pytorch dataloader


I am using librosa to convert some greeting phases as audio files to an STFT spectrogram. I the used the ImageFolder command to load and covert the images to tensors using transforms.ToTensor() as transforms


TRAIN_DATA_PATH = ‘…\librosa\images’
train_data = torchvision.datasets.ImageFolder(root=TRAIN_DATA_PATH, transform=TRANSFORM_IMG)
train_data_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)


There are a few things I want to do but did not find clear cut explanations on them.

I want to visualize each of the images from the data loader. Can someone explain with some code about how to do this? I tried looking in this forum but found a few ways. Can someone give me one standard way with an example?

I tried this from the udemy pytorch course but got an error

dataiter = iter(train_data_loader)
images, labels =

for idx in np.arange(2):
ax = fig.add_subplot(2,20/2,idx+1, xticks=[],yticks=[])

Error: TypeError: Invalid shape (3, 256, 256) for image data