Fix integration slow tests (#10670)

* PoC

* Fix slow tests for the PT1.8 Embedding problem
This commit is contained in:
Sylvain Gugger
2021-03-11 13:43:53 -05:00
committed by GitHub
parent 3ab6820370
commit fda703a553
9 changed files with 47 additions and 55 deletions

View File

@@ -291,13 +291,14 @@ class AlbertModelTest(ModelTesterMixin, unittest.TestCase):
class AlbertModelIntegrationTest(unittest.TestCase):
@slow
def test_inference_no_head_absolute_embedding(self):
model = AlbertForPreTraining.from_pretrained("albert-base-v2")
model = AlbertModel.from_pretrained("albert-base-v2")
input_ids = torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
output = model(input_ids)[0]
expected_shape = torch.Size((1, 11, 30000))
attention_mask = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
output = model(input_ids, attention_mask=attention_mask)[0]
expected_shape = torch.Size((1, 11, 768))
self.assertEqual(output.shape, expected_shape)
expected_slice = torch.tensor(
[[[4.6061, 0.7321, -1.7725], [4.6061, 0.7323, -1.7727], [4.6061, 0.7323, -1.7727]]]
[[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]
)
self.assertTrue(torch.allclose(output[:, :3, :3], expected_slice, atol=1e-4))
self.assertTrue(torch.allclose(output[:, 1:4, 1:4], expected_slice, atol=1e-4))