I’m not entirely sure what the mask is for. Is it because you only want to compute the loss over certain elements of your output sequence (i.e. not on the padding)?
If so, you can pass an argument ignore_index=0 to the call to F.cross_entropy(). That way elements with index 0 will not contribute to the loss. See the docs 1 for details.
You also seem to have put your entire training loop inside your model’s forward() function. I’d recommend making your model compute just the mapping from input sequence to output sequence, and taking everything else out.