NLLLoss is directly computed on log of Softmax and with some high learning rates you can get
unstable trainings/inf loss and RuntimeWarning because of div by 0 in nlls = -np.log(probs_targets)

example with the same lr = 0.1 in numpy and pytorch implementations:

To get it more in line with pytorch I further simplify the log of the softmax in order to use the LogSumExp trick.
You had already done exp_logits = np.exp(logits - logits_max), I just take the log of the sum of that.

I get the actual probabilities back for the cache by taking the exponential of log_probs

This doesn't have much impact on this result tbh and could simply be left as an exercise idea.