Solved by ptrblck in post #2 
              
              
                Images stored as numpy arrays are expected in the channels-last memory layout. 
For a grayscale image you could thus pass X as [256, 256] or [256, 256, 1], which will return the expected output. 
In the current format, torchvision assumes that your image has a height of 1, a width of 256, and 256 ch…