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_electra.py
+++ b/tests/test_modeling_electra.py
@@ -39,7 +39,15 @@ if is_torch_available():
 class ElectraModelTest(ModelTesterMixin, unittest.TestCase):

    all_model_classes = (
-        (ElectraModel, ElectraForMaskedLM, ElectraForTokenClassification,) if is_torch_available() else ()
+        (
+            ElectraModel,
+            ElectraForPreTraining,
+            ElectraForMaskedLM,
+            ElectraForTokenClassification,
+            ElectraForSequenceClassification,
+        )
+        if is_torch_available()
+        else ()
    )

    class ElectraModelTester(object):