SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix integration slow tests (#10670)

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_modeling_squeezebert.py
View File

@@ -278,9 +278,9 @@ class SqueezeBertModelIntegrationTest(unittest.TestCase):
def test_inference_classification_head(self):
model = SqueezeBertForSequenceClassification.from_pretrained("squeezebert/squeezebert-mnli")
input_ids = torch.tensor([[0, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]])
input_ids = torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]])
output = model(input_ids)[0]
expected_shape = torch.Size((1, 3))
self.assertEqual(output.shape, expected_shape)
expected_tensor = torch.tensor([[0.5075, 0.0682, -0.5881]])
expected_tensor = torch.tensor([[0.6401, -0.0349, -0.6041]])
self.assertTrue(torch.allclose(output, expected_tensor, atol=1e-4))