As your script is quite complicated, you could try to build PyTorch from source and try out the anomaly detection, which will try to get the method causing the NANs.
You’ll find the build instructions here.
Let me know, if you encounter any problems.
Alternatively, you could create an executable code snippet and I could try to run it on my machine.