Function detect_nan_hook

training/run.py:84–88

(grad, name)

    )

    def detect_nan_hook(grad, name):
        if torch.isnan(grad).any():
            print(f"NaN detected in gradients of {name}!")
            print(f"Gradient values: {grad}")
            exit()

    # 注册钩子到每个参数
    for name, param in model.named_parameters():