Hi, I am trying to implement this 42 paper.
My implementation of the paper is here 6 for any information about the architecture.
I am calculating the loss manually by taking a negative log of the probability. (Probability of the target sequence is calculated by the equation defined in the paper)
The error comes after few backward passes :
Warning: Traceback of forward call that caused the error:
  File "main.py", line 455, in <module>
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
  File "main.py", line 381, in train
    output = model(src, trg)
  File "/home/shreyansh/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "main.py", line 346, in forward
    output = self.decoder(trg, hidden, encoder_outputs)
  File "/home/shreyansh/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "main.py", line 308, in forward
    probabilities = self.stable_softmax(prediction)
  File "main.py", line 250, in stable_softmax
    numerator = torch.exp(z)
 (print_stack at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:57)
Traceback (most recent call last):
  File "main.py", line 455, in <module>
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
  File "main.py", line 385, in train
    loss.backward()
  File "/home/shreyansh/anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/shreyansh/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'ExpBackward' returned nan values in its 0th output.

The error is coming while calculating gradients of softmax.
  def stable_softmax(self,x):
    z = x - torch.max(x,dim=1,keepdim=True).values
    numerator = torch.exp(z)
    denominator = torch.sum(numerator, dim=1, keepdims=True)
    softmax = numerator / denominator
    return softmax

I am running it on GPU.
torch.isnan(x).any() gives False tensor.
torch.isnan(z).any() gives False tensor.
So, that implies there are no NaNs in the forward pass. I am unable to figure out how NaN is coming in the backward pass.