Hello
I am using pytorch for build MLP neural network for multiclass classification.
epochs = 100 # to have 100 training epochs
for epoch in range(epochs):
batch_loss = 0
for data in train_loader:
    input, output = data
    
    #it is required to zero out the gradients so that we update parameter update correctly. Otherwise the gradient may not point towards the indetended direction than towards the minimum 
    optimizer.zero_grad() 

    predicted = network.forward(input) # a batch of 10 will be applied. neural network will apply the sum function and the sigmoid activation function and it will do prediction on that basis
    print(predicted.shape)
    print("output",output.shape)
   
    #_,pred_label = torch.max(predicted, dim = 1)
    #correct = (pred_label == output).float()

    #loss = criterion(y_pred.squeeze(), y_train)
    loss = loss(predicted.squeeze(), output) #predicted.shape (10, 5) , output.shape - 10
    loss.backward() #to apply the back propogation algorithm
    optimizer.step() #to update the weights

    batch_loss += loss.item() #error which is calculated for each of the batch

#to print the loss for each epoch 
print('Epoch: ' + str(epoch + 1) + ' has a average loss of : ' + str(batch_loss / len(train_loader)))

on running the above code, I am getting following error -
RuntimeError                              Traceback (most recent call last)
 in 
17
18         #loss = criterion(y_pred.squeeze(), y_train)
—> 19         loss = loss(predicted.squeeze(), output) #to calculate the mean squared error defined in above cell.
20         loss.backward() #to apply the back propogation algorithm
21         optimizer.step() #to update the weights
C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
887             result = self._slow_forward(*input, **kwargs)
888         else:
 → 889             result = self.forward(*input, **kwargs)
890         for hook in itertools.chain(
891                 _global_forward_hooks.values(),
C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\loss.py in forward(self, input, target)
526
527     def forward(self, input: Tensor, target: Tensor) → Tensor:
 → 528         return F.mse_loss(input, target, reduction=self.reduction)
529
530
C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\functional.py in mse_loss(input, target, size_average, reduce, reduction)
2923         reduction = _Reduction.legacy_get_string(size_average, reduce)
2924
 → 2925     expanded_input, expanded_target = torch.broadcast_tensors(input, target)
2926     return torch._C._nn.mse_loss(expanded_input, expanded_target, _Reduction.get_enum(reduction))
2927
C:\ProgramData\Anaconda3\lib\site-packages\torch\functional.py in broadcast_tensors(*tensors)
72     if has_torch_function(tensors):
73         return handle_torch_function(broadcast_tensors, tensors, *tensors)
—> 74     return _VF.broadcast_tensors(tensors)  # type: ignore
75
76
RuntimeError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 1
How can I rectify it?
Thanks