Hi Amin!
Most typically for semantic segmentation you will classify each pixel in
your image  into one of nClass classes (including what’s usually called
a “background” class).
The input to your model will have shape [nBatch, nChannel, H, W]
(where the batch size, nBatch, is not built into the model, and the
model is (usually) agnostic to it).  For an RGB image, nChannel would
be 3.
The final layer of your U-Net model would typically be a Conv2d with
out_channels = nClass and the output of the model would have
shape [nBatch, nClass, H, W].
In this situation you would typically use (or at least start with)
CrossEntropyLoss as the loss function, and your ground-truth
labels, the target passed to CrossEntropyLoss, will have shape
[nBatch, H, W] (no nClass dimension), and consist of integer
class labels that run from 0 to nClass - 1.
Best.
K. Frank