@ptrblck thanks so much for reaching out! Your answer certainly solved the dimensional weirdness.
Unfortunately there is still an issue with the grayscale. A sample image test_img[0] has the shape (1, 28, 28). If i try plt.imshow(test_img[0].permute(1, 2, 0) that is one dimension too many for imshow(). If I try plt.imshow(test_img[0].squeeze() I get the following picture. That sample image seem neither grayscaled nor appropriately resized. Any ideas?