banikr:

Is it possible to find out what becomes nan first?


Yes, that was the suggestion in my previous post. You could add print statements in the forward method and check, which activation gets these invalid values first to further isolate it.
Also, if the invalid values are created in the backward pass, you could use torch.autograd.set_detect_anomaly(True) and rerun the code.