Fix the tests for Electra (#6284)

* Fix the tests for Electra * Apply style
2020-08-07 15:30:57 +02:00
parent 6ba540b747
commit 0e36e51515
2 changed files with 11 additions and 3 deletions
--- a/tests/test_modeling_tf_electra.py
+++ b/tests/test_modeling_tf_electra.py
@@ -63,6 +63,7 @@ class TFElectraModelTester:
        self.num_labels = 3
        self.num_choices = 4
        self.scope = None
+        self.embedding_size = 128

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
@@ -194,7 +195,14 @@ class TFElectraModelTester:
 class TFElectraModelTest(TFModelTesterMixin, unittest.TestCase):

    all_model_classes = (
-        (TFElectraModel, TFElectraForMaskedLM, TFElectraForPreTraining, TFElectraForTokenClassification,)
+        (
+            TFElectraModel,
+            TFElectraForMaskedLM,
+            TFElectraForPreTraining,
+            TFElectraForTokenClassification,
+            TFElectraForMultipleChoice,
+            TFElectraForSequenceClassification,
+        )
        if is_tf_available()
        else ()
    )