def METHOD_NAME(self):
    x = torch.randn(2, 3, 4)
    linear = Linear2D(4, 5)
    old_linear = nn.Linear(4, 5)
    old_linear.load_state_dict(linear.state_dict())
    y = linear(x)
    self.assertEqual(y.size(), (2, 3, 5))
    y1 = old_linear(x)
    self.assertEqual(y1.size(), (2, 3, 5))
    torch.testing.assert_close(y, y1)
    Y = torch.randn_like(y)
    loss = torch.nn.functional.mse_loss(y, Y)
    loss1 = torch.nn.functional.mse_loss(y1, Y)
    loss.backward()
    loss1.backward()
    name2grad = {}
    for name, p in linear.named_parameters():
        name2grad[name] = p.grad
    for name, p in old_linear.named_parameters():
        torch.testing.assert_close(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)