Thank you for your prompt response!  I am also wondering about this line:



 ptrblck:

logp = F.log_softmax(logits)


I understand that it is part of the loss formula, but I thought Softmax is computing the probabilities along one dimension. In the 2D-case this would mean along the channel dimension, right? Or are you computing the softmax along HxW? For me, it does not work: The log_softmax returns a tensor full of zeros with shape (B, C, H, W), because I have a pixel-wise two-class classification problem: background vs. foreground. Therefore the number of output channels and the number of target channels equals 1, in my case.
Am I missing something? I appreciate your help! 
def lossWithWeightmap(self, logit_output, target, weight_map):
    #logit_output has shape (B, C, H, W)
	logSoftmaxOutput = nn.functional.log_softmax(logit_output)

    #target has shape (B, 1, H, W)	
    target = target.type(torch.long)
	logSoftmaxOutput = logSoftmaxOutput.gather(1, target.view(self.batchSize, 1, self.imgSize, self.imgSize))
	weightedOutput = (logSoftmaxOutput* weight_map).view(self.batchSize, -1)
	weightedLoss = weightedOutput.sum(1) / weight_map.view(self.batchSize, -1).sum(1)
	weightedLoss = -1.0*weightedLoss.mean()
	return weightedLoss