Solved by KFrank in post #2 
              
              
                Hi Danish! 

Based on this I speculate that you are performing binary segmentation 
of an image.  That is, for each pixel in your sample image, you want to 
predict whether it is “background” or “foreground.” 
If so, you should use be using BCEWithLogitsLoss. 
Let’s say that your sample images have …