While the checkpointing code looks a bit complicated, I cannot see any obvious error.
Could you load two models, one on the CPU and the other one on the GPU, and compare all parameters of these models?
Also, I assume your data pipeline hasn’t changed between the CPU and GPU runs?
Could you nevertheless check the inputs for NaN and Inf values?