Loss_res is NaN

[arid29]

Hi, As I was trying to reproduce the code, I found that Loss_res was somehow NaN which was making it not-differentiable, thus breaking the code. Can you please help me with the same?

Thanks

[MickeyLLG]

I revised the code and this works for me.

class ResLoss(nn.Module):

def __init__(self, reduction='mean'):
    super(ResLoss, self).__init__()
    self.criterion = nn.KLDivLoss(reduction='none')
    self.reduction = reduction

def forward(self, output, target, target_weight=None):
    B, K, _, _ = output.shape
    heatmaps_pred = output.reshape((B, K, -1))
    mark_i = torch.argmax(heatmaps_pred, dim=-1)

    heatmaps_gt = target.reshape((B, K, -1))
    mark_j = torch.argmax(heatmaps_gt, dim=-1)

    mask = torch.ones_like(heatmaps_pred)
    mask[:, :, mark_i] = 0
    mask[:, :, mark_j] = 0
    heatmaps_pred = heatmaps_pred * mask
    heatmaps_gt = heatmaps_gt * mask

    heatmaps_pred = F.log_softmax(heatmaps_pred, dim=-1)
    heatmaps_gt = F.softmax(heatmaps_gt, dim=-1)
    loss = F.kl_div(heatmaps_pred, heatmaps_gt, reduction=self.reduction)
    return loss