(pytorch / mse) How can I change the shape of tensor?
Problem definition:

I have to use MSELoss function to define the loss to classification problem. Therefore it keeps saying the error message regarding the shape of tensor. 

Entire error message:


  torch.Size([32, 10]) torch.Size([32])
  --------------------------------------------------------------------------- RuntimeError                              Traceback (most recent call
  last)  in 
       53     output = model.forward(images)
       54     print(output.shape, labels.shape)
  ---> 55     loss = criterion(output, labels)
       56     loss.backward()
       57     optimizer.step()
  
  /opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in
  call(self, *input, **kwargs)
      530             result = self._slow_forward(*input, **kwargs)
      531         else:
  --> 532             result = self.forward(*input, **kwargs)
      533         for hook in self._forward_hooks.values():
      534             hook_result = hook(self, input, result)
  
  /opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py in
  forward(self, input, target)
      429 
      430     def forward(self, input, target):
  --> 431         return F.mse_loss(input, target, reduction=self.reduction)
      432 
      433 
  
  /opt/conda/lib/python3.7/site-packages/torch/nn/functional.py in
  mse_loss(input, target, size_average, reduce, reduction)    2213
  ret = torch.mean(ret) if reduction == 'mean' else torch.sum(ret)
  2214     else:
  -> 2215         expanded_input, expanded_target = torch.broadcast_tensors(input, target)    2216         ret =
  torch._C._nn.mse_loss(expanded_input, expanded_target,
  _Reduction.get_enum(reduction))    2217     return ret
  
  /opt/conda/lib/python3.7/site-packages/torch/functional.py in
  broadcast_tensors(*tensors)
       50                 [0, 1, 2]])
       51     """
  ---> 52     return torch._C._VariableFunctions.broadcast_tensors(tensors)
       53 
       54 
  
  > RuntimeError: The size of tensor a (10) must match the size of tensor
  b (32) at non-singleton dimension 1


How can I reshape the tensor, and which tensor (output or labels) should I change to calculate the loss?

Entire code is attached below. 

import numpy as np
import torch

# Loading the Fashion-MNIST dataset
from torchvision import datasets, transforms

# Get GPU Device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize((0.5,), (0.5,))])
# Download and load the training data
trainset = datasets.FashionMNIST('MNIST_data/', download = True, train = True, transform = transform)
testset = datasets.FashionMNIST('MNIST_data/', download = True, train = False, transform = transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size = 32, shuffle = True, num_workers=4)
testloader = torch.utils.data.DataLoader(testset, batch_size = 32, shuffle = True, num_workers=4)

# Examine a sample
dataiter = iter(trainloader)
images, labels = dataiter.next()

# Define the network architecture
from torch import nn, optim
import torch.nn.functional as F

model = nn.Sequential(nn.Linear(784, 128),
                      nn.ReLU(),
                      nn.Linear(128, 10),
                      nn.LogSoftmax(dim = 1))
model.to(device)

# Define the loss
criterion = nn.MSELoss()

# Define the optimizer
optimizer = optim.Adam(model.parameters(), lr = 0.001)

# Define the epochs
epochs = 5

train_losses, test_losses = [], []

for e in range(epochs):
  running_loss = 0
  for images, labels in trainloader:
    # Flatten Fashion-MNIST images into a 784 long vector
    images = images.to(device)
    labels = labels.to(device)
    images = images.view(images.shape[0], -1)

    # Training pass
    optimizer.zero_grad()
    output = model.forward(images)
    print(output.shape, labels.shape)
    loss = criterion(output, labels)
    loss.backward()
    optimizer.step()

    running_loss += loss.item()
  else:
    test_loss = 0
    accuracy = 0

    # Turn off gradients for validation, saves memory and computation
    with torch.no_grad():
      # Set the model to evaluation mode
      model.eval()

      # Validation pass
      for images, labels in testloader:
        images = images.to(device)
        labels = labels.to(device)
        images = images.view(images.shape[0], -1)
        ps = model(images)
        test_loss += criterion(ps, labels)
        top_p, top_class = ps.topk(1, dim = 1)
        equals = top_class == labels.view(*top_class.shape)
        accuracy += torch.mean(equals.type(torch.FloatTensor))

    model.train()

    print("Epoch: {}/{}..".format(e+1, epochs),
          "Training loss: {:.3f}..".format(running_loss/len(trainloader)),
          "Test loss: {:.3f}..".format(test_loss/len(testloader)),
          "Test Accuracy: {:.3f}".format(accuracy/len(testloader)))


From the output you print before it error, torch.Size([32, 10]) torch.Size([32]).

The left one is what the model gives you and the right one is from trainloader, normally you use this for something like nn.CrossEntropyLoss.

And from the full error log, the error is from this line 

loss = criterion(output, labels)


The way to make this work is called One-hot Encoding, if it's me for sake of my laziness I'll write it like this.

ones = torch.sparse.torch.eye(10).to(device)  # number of class class
labels = ones.index_select(0, labels)

