I assume you are working on a multi-class segmentation use case, since you are using nn.CrossEntropyLoss.
If that’s the case, your output should have the shape [batch_size, nb_classes, height, width].
While the number of dimensions is correct, it seems you are only dealing with a single class.
Also, the target is expected to have the shape [batch_size, height, width] and contain the class indices in the range [0, nb_classes-1], while your target has the shape [batch_size, 1, 2].
Also, if you are working on a binary segmentation use case, you should either increase the number of output channels of the last conv to 2 or use nn.BCEWithLogitsLoss instead.