U-Net can't perfectly segment image

vision
27

so i just need to set it like this?

criterion = nn.BCEWithLogitsLoss(pos_weight[torch.isinf(pos_weight)|torch.isnan(pos_weight)] = 124.26)