Hi Patrice, For a segmentation model, i would suggest you to use a soft dice loss rather than using weighted nn.BCEWithLogitsLoss(). I don’t think you would see much improvments even with a weight cross entropy loss. You can implement soft dice loss like this:
import torch.nn as nn
import torch

class SoftDiceLoss(nn.Module):
    def __init__(self, weight=None, size_average=True):
        super(SoftDiceLoss, self).__init__()

    def forward(self, logits, targets):
        smooth = 1
        num = targets.size(0)
        """
       I am assuming the model does not have sigmoid layer in the end. if that is the case, change torch.sigmoid(logits) to simply logits
        """
        probs = torch.sigmoid(logits)
        m1 = probs.view(num, -1)
        m2 = targets.view(num, -1)
        intersection = (m1 * m2)

        score = 2. * (intersection.sum(1) + smooth) / (m1.sum(1) + m2.sum(1) + smooth)
        score = 1 - score.sum() / num
        return score

criterion = SoftDiceLoss().to(device) # to put on cuda or cpu
# In the training loop
output = model(images)
loss = criterion(output, gt_masks)
# Then do the backwards and optimizer step