I have a CNN model to classify the MNIST digits problem. The model has 10 output nodes, used CrossEntropyLoss as a loss function and Adam optimizer. 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),                
        )
        # fully connected layer, output 10 classes
        self.out = nn.Linear(32 * 7 * 7, 10)
        # self.softmax = torch.nn.Softmax(dim=1)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        # flatten the output of conv2 to (batch_size, 32 * 7 * 7)
        x = x.view(x.size(0), -1)       
        output = self.out(x)
        # output = self.softmax(output)
        return output, x    # return x for visualization

As, we can use CNN to predict the house price prediction from images. Now, I want to consider this classification problem as a regression problem. I want to predict the digit as a number (whatever the model predicts, i.e, 1.0, 0.5, 8.5, etc). To do this I need only 1 output neuron instead of 10.
If I change this line for the expected output (1 regression output), I am getting an error (as the target labels are from 0 to 9).
# fully connected layer, 1 output
self.out = nn.Linear(32 * 7 * 7, 1)

Error
IndexError: Target 2 is out of bounds.

Could you tell me what I have to do, if I want to predict only one number as digits? Moreover, do I need to change the loss function to MSELoss and optimizer to SGD?
The data-loader
def data_loaders():
    train_data = datasets.MNIST(
        root = 'data',
        train = True,                         
        transform = transforms.ToTensor(), 
        download = True,            
    )
    test_data = datasets.MNIST(
        root = 'data', 
        train = False, 
        transform = transforms.ToTensor()
    )

    train_test_loaders = {
        'train' : torch.utils.data.DataLoader(train_data, 
                                            batch_size=100, 
                                            shuffle=True, 
                                            num_workers=1),
        
        'test'  : torch.utils.data.DataLoader(test_data, 
                                            batch_size=100, 
                                            shuffle=True, 
                                            num_workers=1),
    }
    return train_test_loaders

Later, at the time of training
def train(NB_EPOCS, model, loaders):
    model.train()
        
    # Train the model
    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
            print(np.shape(b_x), np.shape(b_y))
            output = model(b_x)[0]               
            loss = criterion(output, b_y)

The shape of the b_x and b_y is
torch.Size([100, 1, 28, 28]) torch.Size([100])