Hi, I am taking the output from my final convolutional transpose layer into a softmax layer and then trying to measure the mse loss with my target. Problem is that my loss is doesn’t decrease and is stuck around the same point. I figured the problem is using the softmax in the last layer.
I need the softmax layer in the last layer because I want to measure the probabilities. Any suggestion on how to overcome this problem would be helpful. Thanks
Ex:
input  = torch.randn(1, 16, 2, 256, 256)
model output  = torch.randn(1, 16, 1, 256, 256)
target = torch.randn(1, 16, 1, 256, 256)
criterion = mse()