Hi Danish!



 danishnazir:

I have a prediction of shape [1,1,136,100]  where Class is one and Batch size is 1.
My Label  is of size [1,136,100]
Label = [[[  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
0,   0,255,…]]]


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 shape [height, width].  Then
an input batch will have shape [nBatch, height, width], and your
prediction will have the same shape, with no class dimension.  So,
using your sizes, both your input batch and prediction will have shape
[1, 136, 100].
Your target (your “Label”) should also have this same shape, and
contain values in the range [0.0, 1.0], inclusive.  For a purely binary
target (e.g., black-and-white), these values can be exactly 0.0 and
1.0.  Note that for a black-and-white segmentation mask you do not
want values of 0 and 255, but rather 0.0 and 1.0.
Lastly, your prediction should be the raw-score logits output by your
final Linear layer, and not be passed through a sigmoid() that would
convert them to probabilities.
Best.
K. Frank