Oh, it’s a little bit hard to identify which layer.
nan can occur for some reasons but mainly it’s oftentimes 0/inf related maths.
For example, in SCAN code (SCAN/model.py at master · kuanghuei/SCAN · GitHub), nan and inf can happen in forward of l1norm and l2norm. So, I think it’s better to investigate where those bad values are generated, for example, by using forward/backward hooks.
I think one implementation of a backward hook with module name is as follows.
Also, forward hook with similar functionalities will be helpful.
# backward hook with module name
def get_backward_hook(module_name: str):
    
    class BackwardHook:
        name: str
            
        def __init__(self, name):
            self.name = name
            
        def __call__(self, module, grad_input, grad_output):
            for i, g_in in enumerate(grad_input):
                print(module_name, torch.any(torch.isnan(g_in)))
                if torch.any(torch.isnan(g_in)):
                    print(f"{module_name}'s {i}th input gradient is nan")
            for i, g_out in enumerate(grad_output):
                if torch.any(torch.isnan(g_out)):
                    print(f"{module_name}'s {i}th output gradient is nan")
                
    return BackwardHook(module_name)

...

model = SomeCoolModel(...)
for name, module in model.named_modules():
    module.register_full_backward_hook(get_backward_hook(name))

Hope this