Thanks @ptrblck for the help, I finally found the issue.

These alpha here denotes probability after projection over target vocabulary, I was implementing this equation as is.
As you can see there is a summation over multiplication of probabilities due to this alpha values were underflowing.
I wasn’t aware of this logsumexp 261 trick that could have been used to avoid underflow of probabilities by computing all probabilities in log scale.
If anyone is interested, go through this link 120 to understand the trick. The implementation can be found here - https://pytorch.org/docs/master/generated/torch.logsumexp.html 261