def METHOD_NAME(self):
    device = torch.cuda.current_device()
    batch, seq_length = 10, 500
    input_dim, output_dim = 128, 666
    linear0 = torch.nn.Linear(input_dim, output_dim)
    x = torch.randn(batch, seq_length, input_dim, device=device)
    mod = AmpFusedDense(input_dim, output_dim)
    mod.to(device)
    linear0.to(device)
    linear0.zero_grad()
    mod.weight.data.copy_(linear0.weight)
    mod.bias.data.copy_(linear0.bias)
    out = mod(x.reshape(batch * seq_length, -1))
    out = out.reshape(batch, seq_length, -1)
    out1 = linear0(x)
    torch.testing.assert_close(out, out1)
    Y = torch.randn_like(out)
    loss = torch.sum(torch.pow(Y - out, 2))
    loss.backward()
    loss = torch.sum(torch.pow(Y - out1, 2))
    loss.backward()
    for (name, p0), p1 in zip(mod.named_parameters(), linear0.parameters()):
        if p0.grad is not None:
            with self.subTest("grad check", name=name):
                torch.testing.assert_allclose(p0.grad, p1.grad)