I want to create a custom loss function for multi-label classification. The idea is to weigh the positive and negative labels differently. For this, I am making use of this custom code implementation.
class WeightedBCEWithLogitLoss(nn.Module):
    def __init__(self, pos_weight, neg_weight):
        super(WeightedBCEWithLogitLoss, self).__init__()
        self.register_buffer('neg_weight', neg_weight)
        self.register_buffer('pos_weight', pos_weight)

    def forward(self, input, target):
        assert input.shape == target.shape, "The loss function received invalid input shapes"
        y_hat = torch.sigmoid(input + 1e-8)
        loss = -1.0 * (self.pos_weight * target * torch.log(y_hat + 1e-6) + self.neg_weight * (1 - target) * torch.log(1 - y_hat + 1e-6))
        # Account for 0 times inf which leads to nan
        loss[torch.isnan(loss)] = 0
        # We average across each of the extra attribute dimensions to generalize it
        loss = loss.mean(dim=1)
        # We use mean reduction for our task
        return loss.mean()

I started getting nan values which I realized happened because of 0 times inf multiplication. I handled it as shown in the figure. Next, I again saw getting inf as the error value and corrected it by adding 1e-6 to the log (I tried with 1e-8 but that still gave me inf error value).
It would be great if someone can take a look and suggest further improvements and rectify any more bugs visible here.