Segmentation tasks can be seen as a pixel-wise classification, if each pixel can only have one single class.
Similar to a vanilla classification we predict the classes for the input, but for each pixel instead of the whole image.
The loss functions like nn.CrossEntropyLoss work for multi-dimensional output.
Here is a small example using a dummy model with just one conv layer:
n_classes = 2
c, h, w = 3, 24, 24

model = nn.Conv2d(c, n_classes, 3, 1, 1)
criterion = nn.CrossEntropyLoss()

x = torch.randn(1, c, h, w)
target = torch.empty(1, h, w, dtype=torch.long).random_(n_classes)
output = model(x)
loss = criterion(output, target)