Hello @moskomule,
Would the following two steps achieve what you want:

Initialize CrossEntropyLoss 5 with a weight of ones except the first class (which gets weight 0). So if you use cuda

loss_class_weights = torch.ones(num_classes).cuda()
loss_class_weights[0] = 0.0
ce_loss = torch.nn.CrossEntropyLoss(weight=loss_class_weights)


Use loss = ce_loss(output.transpose(0,2,3,1)).continguous().view(-1,1), target.view(-1))


The only thing you would need to take care of is the scaling of the loss - I think that by default it would now average over target.size() instead of (target!=0).sum() - so you would need to scale by the quotient.
Best regards
Thomas
Edit: The transpose above is incorrect, I had numpy’s transpose in mind. See below for what I actually wanted.