Sorry, my bad…
So you get 0 logit and thus infty score and feed it in the nll. This gives you 0 grad in the backward for this entry, which then gets multiplied by 1/0=inf (the derivative of log at 0), and so 0*inf=NaN. This is the curse of backpropagation because next it’ll be multiplied by 0 again.
With log_softmax you would not run into the problem, because it can operate entirely in log-space (i.e. you have logsoftmax(x) = x - logsumexp(x) and so you avoid the derivative of log) and so it can tell that you only have 0.
But is your loss actually a vector of log probabilities? I would seem the logsumexp is inf due to the sign flip. It would need to be log probabilities for mathematical meaning.
Best regards
Thomas