I have a regression model with multiple outputs, so I have a custom mse function

```
def custom_mse(predicted, target):
total_mse = 0
for i in range(target.shape[1]):
total_mse+=nn.MSELoss()(predicted[i], target[i])
return total_mse
```

the purpose is to add up all the mse for each individual output, and return the total as the loss. This works well enough, but is there a way to get rid of the for loop and vectorize the `total_mse`

calculation?

`nn.MSELoss`

accepts batches of your model output and targets.

I’m a bit confused about the indexing. You are iterating using `target.shape[1]`

, but are indexing both tensors in dim0. Is this a typo?

The data has `target.shape`

= `(n_samples, k)`

where k > 1. I’m trying to have a model predict the soft pseudolabels generated by a teacher model, so the targets are the predicted logits for each class, and I’m using the sum of MSE for each class as the loss function

I’m a bit confused about the indexing. You are iterating using `target.shape[1]`

, but are indexing both tensors in dim0. Is this a typo?

Yes, that appears to be an error in the code, it should be indexing `[:,i]`

I did some further digging, it looks like I had an error in my original implementation of `nn.MSELoss`

. Fixing that seemed to give me the results I’m looking for.