Aha! I knew it had to be something silly. On my local machine, I imported the standard MNIST dataset in which the 28x28 pixel arrays are already flattened to 1d arrays of size 784, so this returned:
x shape =  torch.Size([50000, 784])
w shape = torch.Size([784, 10])

as expected. In contrast, since as far as I’m aware Colab doesn’t support file i/o directly to/from any local drive, I imported MNIST from keras instead (see above), which apparently does not flatten the arrays, so this returned:
x shape =  torch.Size([50000, 28, 28])
w shape = torch.Size([784, 10])

Hence the error: apparently different MNIST sources store the data in different formats (good to know). Thanks for narrowing it down!
Obviously, one option is to simply flatten the elements of the keras dataset after loading; but it turns out the keras data isn’t properly normalized either (in contrast to the pickled MNIST data on my local machine, in which the pixel values lie between 0 and 1), and I wasn’t sure whether one should normalize each image relative to the largest pixel value therein, or to the largest pixel value in the entire dataset. But since I had to link my Google Drive to write parameters etc. for later analysis anyway, I just uploaded the flattened-and-normalized MNIST dataset there and imported it as usual. Just FYI, in case anyone else runs into similar issues using MNIST in Colab.
Thanks again @ptrblck!