DataLoader vs convolution tensor shape

1

Hi there, I’m probably missing something simple, but can’t get it though.
The goal is to test the MNIST example with a custom image dataset:

from torchvision.datasets import ImageFolder
from torchvision.transforms import ToTensor
data = ImageFolder(root='PytorchTestImgDir', transform=ToTensor()) 
print(data.classes) 
from torch.utils.data import DataLoader
train_loader = DataLoader(data, batch_size=1)

but the result is a:

RuntimeError: Need input of dimension 4 and input.size[1] == 1 but got input to be of shape: [1 x 3 x 138 x 138] at /Users/soumith/anaconda/conda-bld/pytorch-0.1.10_1488750409207/work/torch/lib/THNN/generic/SpatialConvolutionMM.c:47

2

Looks like the MNIST example’s model expects one color channel (MNIST is black and white) but the images you’re providing are being loaded with three (RGB) channels.

1 Like
3

My blunder, thanks a lot for figuring it out!

4

nst as bx or not, no right wrong