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 as F
from torch import optim
from torchvision import transforms
_tasks = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
from torchvision.datasets import MNIST
mnist = MNIST(“data”, download=True, train=True, transform=_tasks)
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
create training and validation split
split = int(0.8 * len(mnist))
index_list = list(range(len(mnist)))
train_idx, valid_idx = index_list[:split], index_list[split:]
create sampler objects using SubsetRandomSampler
tr_sampler = SubsetRandomSampler(train_idx)
val_sampler = SubsetRandomSampler(valid_idx)
create iterator objects for train and valid datasets
trainloader = DataLoader(mnist, batch_size=256, sampler=tr_sampler)
validloader = DataLoader(mnist, batch_size=256, sampler=val_sampler)
class Model(nn.Module):
def init(self):
super().init()
self.hidden = nn.Linear(784, 128)
self.output = nn.Linear(128, 10)
def forward(self, x):
    x = self.hidden(x)
    x = F.sigmoid(x)
    x = self.output(x)
    return x

model = Model()
loss_function = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay= 1e-6, momentum = 0.9, nesterov = True)
for epoch in range(1, 11): ## run the model for 10 epochs
train_loss, valid_loss = [], []
training part
model.train()
for data, target in trainloader:
optimizer.zero_grad()
1. forward propagation
output = model(data)
2. loss calculation
loss = loss_function(output, target)
3. backward propagation
loss.backward()
4. weight optimization
optimizer.step()
train_loss.append(loss.item())
evaluation part
model.eval()
for data, target in validloader:
output = model(data)
loss = loss_function(output, target)
valid_loss.append(loss.item())
Executing this I am getting the following error :

RuntimeError                              Traceback (most recent call last)
 in ()
----> 1 output = model(data)
2
3 ## 2. loss calculation
4 loss = loss_function(output, target)
5
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in call(self, *input, **kwargs)
487             result = self._slow_forward(*input, **kwargs)
488         else:
–> 489             result = self.forward(*input, **kwargs)
490         for hook in self._forward_hooks.values():
491             hook_result = hook(self, input, result)
 in forward(self, x)
6
7     def forward(self, x):
----> 8         x = self.hidden(x)
9         x = nn.AdaptiveAvgPool(x)
10         x = F.sigmoid(x)
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in call(self, *input, **kwargs)
487             result = self._slow_forward(*input, **kwargs)
488         else:
–> 489             result = self.forward(*input, **kwargs)
490         for hook in self._forward_hooks.values():
491             hook_result = hook(self, input, result)
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/linear.py in forward(self, input)
65     @weak_script_method
66     def forward(self, input):
—> 67         return F.linear(input, self.weight, self.bias)
68
69     def extra_repr(self):
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in linear(input, weight, bias)
1352         ret = torch.addmm(torch.jit._unwrap_optional(bias), input, weight.t())
1353     else:
-> 1354         output = input.matmul(weight.t())
1355         if bias is not None:
1356             output += torch.jit._unwrap_optional(bias)
RuntimeError: size mismatch, m1: [3584 x 28], m2: [784 x 128] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:940