RuntimeError: The size of tensor a (256) must match the size of tensor b (32) at non-singleton dimension 3

2

Hi,

Your custom dataset does not have images with same shape. You can add a Resize as a transform in your dataloader so all images can be stacked to each other when dataloader tries to create a batch.

You can use following transform instead:

tfs = transforms.Compose([
            transforms.Resize(img_size),
            transforms.ToTensor(),
            # normalize,
        ])

Bests