Add more models to common tests (#4910)

2020-06-10 13:19:53 -04:00
parent 3b3619a327
commit 4e10acb3e5
9 changed files with 59 additions and 9 deletions
--- a/tests/test_modeling_distilbert.py
+++ b/tests/test_modeling_distilbert.py
@@ -38,7 +38,13 @@ if is_torch_available():
 class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):

    all_model_classes = (
-        (DistilBertModel, DistilBertForMaskedLM, DistilBertForQuestionAnswering, DistilBertForSequenceClassification)
+        (
+            DistilBertModel,
+            DistilBertForMaskedLM,
+            DistilBertForQuestionAnswering,
+            DistilBertForSequenceClassification,
+            DistilBertForTokenClassification,
+        )
        if is_torch_available()
        else None
    )