From 31c799a0c9ee635ead8957523c15f89cf6ebdae0 Mon Sep 17 00:00:00 2001 From: Funtowicz Morgan Date: Mon, 18 May 2020 13:24:41 +0000 Subject: [PATCH] Tag onnx export tests as slow (#4432) --- tests/test_onnx.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_onnx.py b/tests/test_onnx.py index 9b7f5a4001..e0565bf718 100644 --- a/tests/test_onnx.py +++ b/tests/test_onnx.py @@ -3,7 +3,7 @@ from os import sep from os.path import dirname, exists from shutil import rmtree -from tests.utils import require_tf, require_torch +from tests.utils import require_tf, require_torch, slow from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline from transformers.convert_graph_to_onnx import convert, ensure_valid_input, infer_shapes @@ -22,11 +22,13 @@ class OnnxExportTestCase(unittest.TestCase): MODEL_TO_TEST = ["bert-base-cased", "gpt2", "roberta-base"] @require_tf + @slow def test_export_tensorflow(self): for model in OnnxExportTestCase.MODEL_TO_TEST: self._test_export(model, "tf", 11) @require_torch + @slow def test_export_pytorch(self): for model in OnnxExportTestCase.MODEL_TO_TEST: self._test_export(model, "pt", 11)