Hello,
If you want to your network to output a probability distribution, it’s better to use a cross entropy loss (in case you have one_hot target, as for classification, or BCE loss in case of 2 classes) or a KL-divergence loss (in case you have a target distribution). Those losses work generally better than mse in such cases (there are theoretical justification of this claim). Little advice, if you want to use cross entropy loss, do not insert a softmax at the end of your model, CrossEntropyLoss implemented on pytorch works directly with input logits for a better numerical precision and stability.
Hope it helps,
Thomas