Transformers: How to use the target mask properly?
The main issue during training is that you haven't right-shifted the input of the decoder, which is probably why you set the diagonals of mask to -inf (when it should be $0$).
Also, just an FYI, although you haven't focused on evaluation/prediction yet, I will explain the evaluation/prediction here as well for completeness, since it works so differently than training, and also since you will need it when generating the graphs for debugging.
Training
Both tgt and tgt_mask should be changed to simulate auto-regressive properties.
You are feeding in tgt as the input to the decoder, where tgt is the ground truth target sequence to predict. tgt should have dimension length(sequence) x batchSize x $\|dict\|$. Additionally, you are feeding in mask where the diagonals are -inf.
Instead, you should do the following:

Add 2 new special tokens to your vocabulary that signals the start and end of decoding, i.e. &lt;START&gt; and &lt;END&gt;. So your vocabulary will need to be extended to 16 events.
Every input to the decoder (training, evaluation, prediciton) should have &lt;START&gt; as the first token. So during training, when you want to feed in tgt to the decoder, you should right-shift tgt by adding in the &lt;START&gt; token at the beginning.

tgt_shifted will now be of dimension length(sequence)+1 x batchSize x $\| dict \|$
tgt_shifted_mask will now be of dimension length(sequence)+1 x length(sequence)+1. Diagonal should be all $0$
the inputs to the decoder should be tgt_shifted, tgt_shifted_mask, and memory
the output of the decoder will have dimension length(sequence)+1 x batchSize x $\| dict\|$, and will not be right-shifted because it will not have &lt;START&gt; as the first word. But it should have &lt;END&gt; as the last word.
Append &lt;END&gt; to the ground truth tgt, and turn it into one-hot encoding, so tgt should have dimension length(sequence)+1 x batchSize x $\|dict\|$. Your loss operator should be some sort of element-wise comparison between tgt and the output of the decoder.



Note that you only run the decoder once per training batch during training, i.e. the decoder simultaneously predicts the logits of all length(sequence) tokens at the same time. On the other hand, during evaluation/prediction, you must run the decoder length(sequence) times, since you can only use the decoder to predict one token at a time.
Evaluation/Prediction
During evaluation/prediction, the model should generate its own sentence entirely from scratch, one word at a time:

out should be initialized as a single token, &lt;START&gt;, and have dimensions $1$  x batchSize x $\|dict\|$.
Do the following in a loop until termination or max output length is reached:

Generate a out_mask for the current version of out. out_mask should be a matrix with dimension length(out) x length(out). Diagonal should be all $0$
Pass out, out_mask, and memory into the decoder layer. The output of the decoder will be a length(out) x batchSize x $d_{model}$ tensor (note that in your case, $d_{model} = \|dict\|$), which indicates tokenIdx x batchIdx x vocabularyIdx
Convert the last time step (e.g. decoderOutput[-1,:,:] if all sequences in your batch have the same length) into probabilities, and take all the events that surpasses some probability threshold and set them to $1$. Set all other events to $0$. This is the new time step that the decoder layer predicts. So add this time step to out to append it's length by $1$.
Repeat the previous 3 steps in a loop until you reach the &lt;END&gt; token, which signals the end of decoding, or until the max output length is reached.



During evaluation, you can calculate the validation loss by using the logits of only the final iteration of the decoder (which should have dimension length(sequence)+1 x batchSize x $\|dict\|$), and compare those logits with the ground truth with the &lt;END&gt; token added at the end (resulting in dimensions length(sequence)+1 x batchSize x $\|dict\|$).