несоответствие размеров, m1: [3584 x 28], м2: [784 x 128] в /pytorch/aten/src/TH/generic/THTensorMath.cpp:940

Ваши входные данные MNIST имеют форму [256, 1, 28, 28], соответствующую [B, C, H, W]. Вам необходимо сгладить входные изображения в один длинный вектор 784 перед подачей его на линейный слой Linear(784, 128), чтобы входное значение стало [256, 784], соответствующее [B, N], где N равно 1x28x28, вашему размеру изображения. Это можно сделать следующим образом:

for data, target in trainloader:

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

        optimizer.zero_grad()
        ...

То же самое необходимо сделать в цикле проверки.