I want to consider MNIST digits as a regression problem (as we can do for the house price prediction).
I used MSELoos and SGD optimizer. The last layer of the CNN model is linear with one neuron. The structure of the model is given below
class CNN(nn.Module):
        def __init__(self):
            super(CNN, self).__init__()
            self.conv1 = nn.Sequential(        
                nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5,stride=1,padding=2,),                              
                nn.ReLU(),                      
                nn.MaxPool2d(kernel_size=2),    
            )
            self.conv2 = nn.Sequential(        
                nn.Conv2d(16, 32, 5, 1, 2),    
                nn.ReLU(),                      
                nn.MaxPool2d(2),                
            )
            self.out = nn.Linear(32 * 7 * 7, 1)
        def forward(self, x):
            x = self.conv1(x)
            x = self.conv2(x)
            x = x.view(x.size(0), -1)      
            output = self.out(x)
            return output, x

Train code snippets
optimizer = optim.SGD(params=model.parameters(), lr=LR)
criterion = nn.MSELoss()

def train(NB_EPOCS, model, loaders):
    model.train()
    total_step = len(loaders['train'])
    for epoch in range(NB_EPOCS):
        for i, (images, labels) in enumerate(loaders['train']):
            b_x = Variable(images)   # batch x
            b_y = Variable(labels)   # batch y
            output, _ = model(b_x)  
            loss = criterion(output, b_y.float())
            optimizer.zero_grad()          
            loss.backward()    
            optimizer.step()            
            if (i+1) % 100 == 0:
                print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                       .format(epoch + 1, NB_EPOCS, i + 1, total_step, loss.item()))

I am getting a warning
UserWarning: Using a target size (torch.Size([100])) that is different to the input size (torch.Size([100, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)

The loss is very high
Epoch [1/10], Step [100/600], Loss: 7.7500
Epoch [1/10], Step [200/600], Loss: 8.3003
Epoch [1/10], Step [300/600], Loss: 8.7280
Epoch [1/10], Step [400/600], Loss: 8.4920
Epoch [1/10], Step [500/600], Loss: 8.5399
Epoch [1/10], Step [600/600], Loss: 8.8300
Epoch [2/10], Step [100/600], Loss: 10.8930
Epoch [2/10], Step [200/600], Loss: 10.0020
Epoch [2/10], Step [300/600], Loss: 7.9896
Epoch [2/10], Step [400/600], Loss: 7.2748
Epoch [2/10], Step [500/600], Loss: 9.4017

However, at the time of testing, the model is predicting all images as 0.
Test Accuracy of the model on the 10000 test images: %.2f 0.16
Prediction Number: [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
Actual Number [3 8 1 8 3 0 3 9 0 9 3 3 3 1 2 9 7 4 9 1 4 7 4 7 4 3 7 3 1 4 1 7 4 1 6 2 5
 0 0 6 8 8 3 2 5 1 6 3 9 3 8 1 4 1 7 7 5 8 3 2 0 4 3 5 9 3 9 4 7 6 0 7 2 3
 9 2 6 7 3 5 6 8 2 3 7 2 6 5 6 3 6 4 5 0 0 7 0 7 6 6]

I think the warning is one of the reasons for the poor performance. Could you tell me how can I resolve the warning and improve the performance?