def test_linear(self):
    self._check_one_layer(nn.Linear(8, 4), torch.randn(16, 8))
    self._check_one_layer(nn.Linear(8, 4), torch.randn(16, 8, 8))
def test_linear(self):
    self._check_one_layer(nn.Linear(8, 4), torch.randn(16, 8))
    self._check_one_layer(nn.Linear(8, 4), torch.randn(16, 8, 8))