Hey Juan,
Thanks for getting back and sharing your suggestions.


To me it seems strange you are using softmax + mse loss.
You are computing some sort of probabilities but then you use a cartesian distance? - The reason for using softmax plus mse loss is a direct implementation of the paper ([2104.08029] T-LEAP: occlusion-robust pose estimation of walking cows using temporal information 2). It is for a pose estimation problem so that I can have confidence maps as the output.


Lastly, just random suggestion, are you using a validation set extracted from the same type of data? If you use 2 different datasets (rather than splitting one into train and Val) you may find this. -  I am using a personal dataset and splitting the same data into an 80:20 train and validation dataset. I am using different videos as one dataset


Despite the scale of the plot is not proper, it seems even the training is not learning much if it can converge in 10 epochs. - I have trained for 80 epochs and the loss remains on the same scale. However, I realized that when I print loss.backward() i get None value as output which also means that the network is not training.