Hello, i’m making some changes on a normal CNN, to be compatible with other model, that i create.
I get the following error:
File “C:/Users/user/.spyder-py3/cnn.py”, line 105, in 
loss = criterion(output,img)
RuntimeError: The size of tensor a (10) must match the size of tensor b (28) at non-singleton dimension 3
This is the CNN:
import time
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as dsets
from torch.autograd import Variable
GPU = 0
def to_img(x):
x = x.view(x.size(0), 1, 28, 28)
return x
train_dataset = dsets.MNIST(root=’./data’,
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = dsets.MNIST(root=’./data’,
train=False,
transform=transforms.ToTensor())
batch_size =251
n_iters = 71800
num_epochs = n_iters / (len(train_dataset) / batch_size)
num_epochs = int(num_epochs)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
class CNNModel(nn.Module):
def init(self):
super(CNNModel, self).init()
    self.cnn1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
    self.relu1 = nn.ReLU()

    self.maxpool1 = nn.MaxPool2d(kernel_size=2)

    self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2)
    self.relu2 = nn.ReLU()

    self.maxpool2 = nn.MaxPool2d(kernel_size=2)

    self.fc1 = nn.Linear(32 * 7 * 7, 10)

def forward(self, x):
    out = self.cnn1(x)
    out = self.relu1(out)
   
    out = self.maxpool1(out)
   
    out = self.cnn2(out)
    out = self.relu2(out)

    out = self.maxpool2(out)

    out = out.view(out.size(0), -1)

    out = self.fc1(out)
   
    return out

model = CNNModel()
if torch.cuda.is_available():
torch.cuda.set_device(GPU)
model.cuda()
print(“cuda OK”)
criterion = nn.MSELoss()
learning_rate = 1e-3
optimizer = torch.optim.Adam(
model.parameters(), lr=learning_rate, weight_decay=0)
iter = 0
start = time.time()
for epoch in range(num_epochs):
for data in train_loader:
img, _ = data
img = Variable(img).cuda()
nova=to_img(img)
    st = time.time()

    optimizer.zero_grad()

    output = model(img)

    loss = criterion(output,img)

    loss.backward()

    optimizer.step()

    iter += 1

       
print('epoch [{}/{}], loss:{:.4f}'
              .format(epoch + 1, num_epochs, loss.item()))

Anyone knows what i am doing wrong?
Thank you,
Mike