@ptrblck

Use nn.BCEWithLogitsLoss as your criterion, make sure your model outputs logits with the shape [batch_size, 1, height, width]

Shouldn’t it be [batch_size, 2, height, width] because my outputs have two classes?

use a target with the same shape containing your labels (0 and 1).

Again, with classes 0 and 1, shouldn’t the number of classes in target be 2?
Edit:
(Taking two classes as : clear class and dark class.)
I just realized what you meant. Correct me if I am wrong but when you say 1 in [batch_size, 1, height, width], BCELosswithlogits is giving a probabilty for that class. A lower value is lower probabilty for that label, and a higher value would be higher value for that label.
However, I still don’t get what is meant when you said:

use a target with the same shape containing your labels (0 and 1)

With you suggestion, won’t I have only one label now (initially I have giving two label something like clear class and dark class)? The prediction from the model for this label should be 0 or 1 (as it is a segmentation model).
Sorry for so many questions.