Hello,
I am trying to implement a BCEWithLogitsLoss version with a different weight for some pixels. The goal is to have good borders in a segmentation task.
The outputs of my net are logits.
My weight map have the same size as my target. (For example [batch_size, height, width]).
My implementation currently looks like this :
def BCE_with_logits_loss_pixel_map(pred, target, pix_map):
    
    # check that pred, target and map have the same dimension
    assert pred.size() == target.size() == pix_map.size(), "Loss Components don't have the same size"
    
    sig_pred = torch.sigmoid(pred)
    unweighted_loss = target * torch.log(sig_pred) + (1 - target) * torch.log(1 - sig_pred)
    weighted_loss = -(pix_map * unweighted_loss)
    mean_weighted_loss = weighted_loss.mean()
    
    return mean_weighted_loss

batch_size = 8
out_channels = 1
W = 128
H = 128
logits = torch.FloatTensor(batch_size, H, W).normal_()
target = torch.LongTensor(batch_size, H, W).random_(0, out_channels)
weights = torch.FloatTensor(batch_size, H, W).random_(1, 3)

loss_val = BCE_with_logits_loss_pixel_map(logits, target, weights)



I would like to know if there are any mistake in my code ?
Thank you.