From 1bf38611a4183f9fd0e0928703ed2264e0305647 Mon Sep 17 00:00:00 2001 From: Lysandre Debut Date: Fri, 6 Aug 2021 18:41:33 +0200 Subject: [PATCH] Put smaller ALBERT model (#13028) --- tests/test_onnx_v2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_onnx_v2.py b/tests/test_onnx_v2.py index 4bf7529514..3e99162d80 100644 --- a/tests/test_onnx_v2.py +++ b/tests/test_onnx_v2.py @@ -195,7 +195,7 @@ if is_torch_available(): ) PYTORCH_EXPORT_DEFAULT_MODELS = { - ("ALBERT", "albert-base-v2", AlbertModel, AlbertConfig, AlbertOnnxConfig), + ("ALBERT", "hf-internal-testing/tiny-albert", AlbertModel, AlbertConfig, AlbertOnnxConfig), ("BART", "facebook/bart-base", BartModel, BartConfig, BartOnnxConfig), ("BERT", "bert-base-cased", BertModel, BertConfig, BertOnnxConfig), ("DistilBERT", "distilbert-base-cased", DistilBertModel, DistilBertConfig, DistilBertOnnxConfig), @@ -229,7 +229,7 @@ class OnnxExportTestCaseV2(TestCase): self.assertTrue(hasattr(onnx_config_class, "default")) tokenizer = AutoTokenizer.from_pretrained(model) - model = model_class(config_class()) + model = model_class(config_class.from_pretrained(model)) onnx_config = onnx_config_class.default(model.config) with NamedTemporaryFile("w") as output: