Fix slow gpu tests lysandre (#4487)

* There is one missing key in BERT

* Correct device for CamemBERT model

* RoBERTa tokenization adding prefix space

* Style
This commit is contained in:
Lysandre Debut
2020-05-20 11:59:45 -04:00
committed by GitHub
parent 6dc52c78d8
commit 14cb5b35fa
3 changed files with 8 additions and 4 deletions

View File

@@ -80,8 +80,9 @@ class AutoModelTest(unittest.TestCase):
model, loading_info = AutoModelForPreTraining.from_pretrained(model_name, output_loading_info=True)
self.assertIsNotNone(model)
self.assertIsInstance(model, BertForPreTraining)
for value in loading_info.values():
self.assertEqual(len(value), 0)
for key, value in loading_info.items():
# Only one value should not be initialized and in the missing keys.
self.assertEqual(len(value), 1 if key == "missing_keys" else 0)
@slow
def test_lmhead_model_from_pretrained(self):