riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-02 03:01:07 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

Fix Yolos ONNX export test (#18606)

Browse Source 
Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar 2022-08-17 10:04:49 +02:00 committed by GitHub
parent 358478e729
commit c99e984657
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/onnx/test_onnx_v2.py
View File

@ -284,6 +284,12 @@ class OnnxExportTestCaseV2(TestCase):
model_class = FeaturesManager.get_model_class_for_feature(feature) model_class = FeaturesManager.get_model_class_for_feature(feature)
config = AutoConfig.from_pretrained(model_name) config = AutoConfig.from_pretrained(model_name)
model = model_class.from_config(config) model = model_class.from_config(config)
# Dynamic axes aren't supported for YOLO-like models. This means they cannot be exported to ONNX on CUDA devices.
# See: https://github.com/ultralytics/yolov5/pull/8378
if model.__class__.__name__.startswith("Yolos") and device != "cpu":
return
onnx_config = onnx_config_class_constructor(model.config) onnx_config = onnx_config_class_constructor(model.config)
if is_torch_available(): if is_torch_available():