Use small ELECTRA model for testing

explosion · Apr 8, 2024 · 2900ee4 · 2900ee4
1 parent c180a12
commit 2900ee4
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/curated_transformers/tests/models/electra/test_encoder.py b/curated_transformers/tests/models/electra/test_encoder.py
@@ -12,11 +12,7 @@
 @pytest.mark.parametrize("with_torch_sdp", [False, True])
 @pytest.mark.parametrize(
     "model_name",
-    [
-        "jonfd/electra-small-nordic",
-        "Maltehb/aelaectra-danish-electra-small-cased",
-        "google/electra-small-discriminator",
-    ],
+    ["explosion-testing/electra-test"],
 )
 def test_encoder(model_name: str, torch_device, with_torch_sdp):
     assert_encoder_output_equals_hf(