Hi,
I think this is the line:



 maito:

patches = transforms.ToPILImage()(patches.squeeze_(0))


The reason is that let’s say you have an image of size [3, 100, 100] and then you extract 4 patches, then patches would be [4, 3, h, w] and this cannot be converted to a single image using ToPILImage as it only accepts 2D or 3D images.
I suggest that you write your own custom transformation by accepting tensors instead of converting to tensor then extracting to patches then again to image. After all, you are again converting it to tensor to use normalize so why not directly implement patch extraction for normalized tensor?
Also, it terms of performance, this approach will be faster as you no longer have to deal with conversion overhead.
Bests