The paper says:

Note that the efficiency of algorithm 1 can, at the expense of clarity, be improved upon by changing the order of computation, e.g. by replacing the last three lines in the loop with the following lines…

So I’m surprised that it should make a noticeable difference, but maybe that is the case. @christianperone would you mind trying the altered version of Adam on your problem? Fingers crossed this might be the solution.
The bit about sparse updates with TensorFlow Adam I would assume don’t matter in this case. I don’t know how PyTorch deals with sparse modules wrt gradient updates, but what TF claims to do sounds like the correct approach.