you sure the shape is ok?



 szahan:

TEST_DATA_PATH = “./data”
TRANSFORM_IMG = transforms.Compose([
transforms.ToTensor(),
])
test_data = torchvision.datasets.ImageFolder(root=TEST_DATA_PATH, transform=TRANSFORM_IMG)
test_data_loader = data.DataLoader(test_data, batch_size=2, shuffle=False)
for i, data in enumerate(test_data_loader):
images,labels = data


this will give you an image of shape [C, H, W]
while plt.imread(img_path) will give you an image of shape [H, W, C]
If you want the image tensor you get from your datalaoder to be of the same shape, then you have to use .permute(1, 2, 0) on your tensor.
On your example this could be done here:
plt.imshow(images[0].permute(1, 2, 0))

Please also keep in mind that .ToTensor() puts your image data into the range [0, 1] by dividing by 255