def test_grad_norm(norm_type, expected):
    """Test utility function for computing the p-norm of individual parameter groups and norm in total."""

    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.param0 = nn.Parameter(torch.rand(3))
            self.param1 = nn.Parameter(torch.rand(2, 1))
            self.param0.grad = torch.tensor([-1.0, 2.0, -3.0])
            self.param1.grad = torch.tensor([[-4.0], [5.0]])
            # param without grad should not contribute to norm
            self.param2 = nn.Parameter(torch.rand(1))

    model = Model()
    norms = grad_norm(model, norm_type)

    assert norms.keys() == expected.keys()
    for k in norms:
        assert norms[k] == pytest.approx(expected[k])
Test utility function for computing the p-norm of individual parameter groups and norm in total.
tests/tests_pytorch/utilities/test_grads.py