Pytorch: ValueError: Too many dimensions: 3 > 2. 9/opt/anaconda3/envs/pytorch/lib/python3.7/site-packages/PIL/Image.py in fromarray(obj, mode)

I am using MNIST data to run my python using pytorch. I like to train only partial data for digits 0 and 1. when I try to print the size of the first image it runs into "ValueError: too many dimensions: 3 > 2).

I am very new to Python. Can someone help to guide figuring root cause. The program runs fine if I don't segment the training data. here is the code snippet

subset_indices = ((train_data.train_labels == 0) + (train_data.train_labels == 1)).nonzero()
train_loader = torch.utils.data.DataLoader(train_data,batch_size=batch_size, shuffle=False,sampler=SubsetRandomSampler(subset_indices))