I think I have an explanation to the differences:
For every decoder step the model in question would compute a linear forward pass on a fixed input.
It could be 0.2 cached this value given that the input is fixed whereas 0.4 recomputed and added the result to the graph at every iteration.
I’m able to run the code without problems with 0.2 and 0.4 now that the linear transform on that fixed input is computed only once and reused during the decoder steps!
Thank you all.