dnrops
/
transformers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add 'with torch.no_grad()' to integration test forward pass (#14808)

This commit is contained in:
Henrik Holm 2021-12-23 10:23:39 +01:00 committed by GitHub
parent d8c09c6541
commit 6a7b9da2ae
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_modeling_albert.py
View File

@ -299,7 +299,8 @@ class AlbertModelIntegrationTest(unittest.TestCase):
model = AlbertModel.from_pretrained("albert-base-v2")
input_ids = torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
attention_mask = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
output = model(input_ids, attention_mask=attention_mask)[0]
with torch.no_grad():
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(