Solved by ptrblck in post #2 
              
              
                matplotlib usually expects the images to have the channel in dim2, so try to permute your image tensor using tensor = tensor.permute(1, 2, 0) before passing it to imshow.