Hey Frank, thanks for the tip. I am trying to replicate a paper and they state:
The last convolutional layer of the decoder is followed by a
transposed-convolutional layer that uses a linear activation followed by a Softmax transformation, resulting in an output
of 17 confidence maps, one per keypoint.
ref : [2104.08029] T-LEAP: occlusion-robust pose estimation of walking cows using temporal information
When i use softmax as the last layer as suggested in the model then my loss stagnates. It would be nice to know how I can use the softmax (before or after, how should my target heatmap look like then) with the mse loss.  An example would be helpful
For now, i am doing something like this:

input = model + torch.nn.Softmax(dim=1) (torch.rand(1,16,2,256,256))
heatmap = torch.rand(1,16,1,256,256)
criterion = torch.nn.MSELoss()
criterion(input, heatmap)