As you’ve already pointed out the issue is created by the missing “class dimension” in the output.
In a multi-class segmentation use case nn.CrossEntropyLoss expects the model output to have the shape [batch_size, nb_classes, height, width], where each “channel” represents an activation maps, which gives the logits (unnormalized probabilities) for each pixel location.
E.g. if you are working on a binary segmentation you could treat it as a 2-class multi-class segmentation using an output in the shape [batch_size, 2, height, width], where the first channel would correspond to the logits of each pixel for class0 and the second channel for class1, respectively.
In your current code you are using a single output channel, so that your model would only be able to predict class0, which makes it a bit useless.
In any case, since you are working on a binary segmentation, you could also keep the single output channel and use nn.BCEWithLogitsLoss alternatively, which would then map a “low” value to class0 and a “high” one to class1. In that case your target should have the same shape as the model output.