I have no real experience with word models. Nevertheless, here are a couple of thoughts.
I can’t see why you need a custom batchNLLLoss. If you pass tensors of the right shape to NLLLoss then it should calculate the loss over an entire batch.
If the loss is really small, then the gradients will probably all be really small and maybe you could increase the learning rate to counterbalance that fact and speed up the learning.