def test_grad_clip_norm_():
    """
    **Overview**:
        Test function of grad clip by norm.
    """
    # Prepare hyper-parameters, batch size=4, action=32
    B, N = 4, 32
    # Generate regression logit and label, in practice, logit is the output of the whole network and requires gradient.
    logit = torch.randn(B, N).requires_grad_(True)
    label = torch.randn(B, N)
    # Define criterion, and compute loss.
    criterion = torch.nn.MSELoss()
    output = criterion(logit, label)
    # Loss backward and compute gradient.
    output.backward()
    # Clip the total norm of gradients.
    grad_clip_norm_(logit, 0.5, 2)
    # Assert the total_norm of the clipped gradient.
    assert isinstance(logit.grad, torch.Tensor)
    grads = logit.grad
    total_norm = torch.norm(torch.stack([torch.norm(g.detach(), 2) for g in grads]), 2)
    assert total_norm < 0.5

    **Overview**:
        Test function of grad clip by norm.
    
chapter7_tricks/grad_clip_norm.py