Hi I am using such a loss:
nll = torch.pow(y_hat - y, 2) * 0.5
e_log_sigma_2 = torch.exp(-log_sigma_2)
u_a = nll * (e_log_sigma_2) + log_sigma_2 * 0.5
loss = u_a.mean()
After training a while. I get some nan and I cannot understand why. Maybe Somebody have face this problem already.
Any suggestion is welcome =)