You could apply nn.LogSoftmax on your output and nn.NLLLoss for your semgmentation task as you would do it on a simple classification task.
nn.NLLLoss will take multi-dimensional input, so you can just calculate your loss as:
output = model(input)  # raw logit output in shape [1, 3, 512, 512]
loss = criterion(F.log_softmax(output,1), target)  # target in shape [1, 512, 512]