"vscode:/vscode.git/clone" did not exist on "8e28c327fc832f7c6acffa2870c27db01658cc43"
Unverified Commit b41cc0b8 authored by Funtowicz Morgan's avatar Funtowicz Morgan Committed by GitHub
Browse files

Fix flaky ONNX tests (#6531)

parent 39c3b1d9
import unittest
from os.path import dirname, exists
from pathlib import Path
from shutil import rmtree
from tempfile import NamedTemporaryFile, TemporaryDirectory
from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
......@@ -72,7 +70,7 @@ class OnnxExportTestCase(unittest.TestCase):
def test_quantize_pytorch(self):
for model in OnnxExportTestCase.MODEL_TO_TEST:
path = self._test_export(model, "pt", 12)
quantized_path = quantize(Path(path))
quantized_path = quantize(path)
# Ensure the actual quantized model is not bigger than the original one
if quantized_path.stat().st_size >= Path(path).stat().st_size:
......@@ -82,11 +80,11 @@ class OnnxExportTestCase(unittest.TestCase):
try:
# Compute path
with TemporaryDirectory() as tempdir:
path = tempdir + "/model.onnx"
path = Path(tempdir).joinpath("model.onnx")
# Remove folder if exists
if exists(dirname(path)):
rmtree(dirname(path))
if path.parent.exists():
path.parent.rmdir()
# Export
convert(framework, model, path, opset, tokenizer)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment