You should use nn.CrossEntropyLoss and can have a look at the docs 2 to see how the loss formula is implemented.
I would recommend to skip the spatial dimensions for now and just take a look at the standard multi-class classification use case, where the output has the shape [batch_size, nb_classes] and the target contains class indices in the range [0, nb_classes-1] with the shape [batch_size].
As you can see in the docs, the output logits can be directly indexed using the target.
For the multi-class segmentation use case, the same formula yould be applied but for each pixel location, since the model output and the target now have additional [..., height, width] dimensions.