The shapes look almost right. For a multi-class segmentation use case you could use nn.CrossEntropyLoss as the criterion, which expects the model output to contain logits in the shape [batch_size, nb_classes, height, width]. The target should have the shape [batch_size, height, width] (remove dim1 in your script via target = target.squeeze(1)) and should contain the class indices in the range [0, nb_classes-1].
Assuming you are dealing with 20 classes, here is a small code example:
output = torch.randn(2, 20, 24, 24, requires_grad=True)
target = torch.randint(0, 20, (2, 24, 24))
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)




 Thummper:

Do I have to calculate an output prediction with output.argmax(1) for input into a loss function?


No, nn.CrossEntropyLoss expects the logits for each class. You could use torch.argmax(output, dim=1) to compute the predictions, where each pixel would contain the the predicted class index.