F.binary_cross_entropy_with_logits expects the model output and target to have the same shape and can be used for a multi-label segmentation (each pixel can belong to zero, one, or multiple classes).
Since you’ve mentioned a multi-class segmentation (each pixel belongs to one class only), you should use nn.CrossEntropyLoss instead (or nn.NLLLoss with F.log_softmax as the last non-linearity).
nn.CrossEntropyLoss expects the model output to have the shape [batch_size, nb_classes, height, width] and the target [batch_size, height, width] containing the class indices in the range [0, nb_classes-1].
Your target seems to have an additional dimension in dim1 and might be a color image in the channels-last format?
If so, you would have to map the colors to class indices.