Ok if they are supposed to be the softmax score then one thing you could do is try just getting the argmax of them using torch.argmax. That would give it the correct shape and it would provide the loss function with the correct label.