Is it possible to get probability(likelihood) for foreground class label?

[c13proto]

Hello Sachin, I am very beginner of Deep Learning , I apologize if my question is stupid.

I trained with my own annotation images (background is 0, foreground is 1 ) and set "--classes=2".
Then I want to know the probability(likelihood) of label 1 for each pixel of test images.
In other word, I want to get heat-map of probability of foreground label.
Is it possible?
And, I want to know whether trying 2 class segmentation on this network is correct approach or not.

[Sachin Mehta]

For computing probability, first apply softmax in output tensor and then take torch.max of the resultant tensor. This will return two values, max probability along the channel and corresponding index value (or class).

See PyTorch documentation for more details

[c13proto]

I see, thank you for your answer. I try to read it!