Hi Frank,
Thanks a lot for your suggestions.
“I would first check that your  target  data doesn’t contain any outlandish values.”
Great advice, turns out I did have some outlandish target values. Removing them and normalizing the data has fixed the issue. However, I am still a bit confused as to why the loss gave values of such different magnitudes just by introducing the numerically stable log-sum-exp.
m, idx = torch.max(value, dim=dim, keepdim=True)
        return m.squeeze(dim) + torch.log(torch.sum(torch.exp(value-m)*(weights.unsqueeze(2)),
                                       dim=dim) + eps)

Does it have something to do with the max we are taking in this function?
Thanks,
Umang