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.