nn.BCEWithLogitsLoss expects the model output to contain logits, so you would have to remove the sigmoid (it’ll be applied internally in a more numerically stable way).
However, this criterion expects the model output and the target to have the same shape, which doesn’t seem to be the case in your workflow.
For a binary segmentation use case, the model should return an output in [batch_size, 1, height, width]. Based on the current shape I assume you are setting the number of output classes to 2 in a multi-class segmentation. If that’s the case, you would have to either change the last layer so that it outputs one class channel or use nn.CrossEntropyLoss alternatively. In the latter case, you would have to remove the channel dimension from the target so that is has the shape [batch_size, height, width] and contains class indices in the range [0, nb_classes-1].