Sorry I didn’t get back to you before now on this.
How does this fit into your previous findings, i.e. neither the model output, nor the parameters or the gradients were having invalid values, but the `optimizer.step()` caused the parameters to become NaNs?
Before I saw the other posts I was trying to reason about this issue, but wasn’t sure how this could be happening.

You are right that part still doesn’t make sense. I see two possible scenarios for what could have happened. 1) I somehow must have made a mistake when I was checking the regular output for NaNs, 2) The way I was checking was to find an error and reload the model to previous epoch, but I didn’t reload the optimizer (ADAM) which I believe has some kind of momentum history in it (not too sure on how ADAM works), so maybe the momentum in the optimizer was still NaN, which corrupted the optimizer step?
(I have since fixed the issue, and have been running the code for 7 days in a bunch of different tests and they are all super stable now, so it works great :))