Size Mismatch, M1: [3584 X 28], M2: [784 X 128] At /pytorch/aten/src/TH/generic/THTensorMath.cpp:940

I have executed the following code and getting the error shown at extreme bottom. I would like to know how to resolve this. thanks import torch.nn as nn import torch.nn.functional

Solution 1:

Your input MNIST data has shape [256, 1, 28, 28] corresponding to [B, C, H, W]. You need to flatten the input images into a single 784 long vector before feeding it to the Linear layer Linear(784, 128) such that the input becomes [256, 784] corresponding to [B, N], where N is 1x28x28, your image size. This can be done as follows:

for data, target in trainloader:

        # Flatten MNIST images into a 784 long vector
        data = data.view(data.shape[0], -1)

        optimizer.zero_grad()
        ...

The same is needed to be done in the validation loop.