Since you’ve now narrowed the first occurrence of the NaNs to the 2nd layer, you could use the hook to check the inputs as well as the module parameters.
The module, input, and output would be pass to the forward hook and I assume you’ve checked the output for NaNs. If so, check the parameters as well as inputs. In case the parameters already contain invalid values, check the gradients for this layer in the previous iteration as they are most likely also containing invalid values.