From f8d873bd8fa8c7cd53ceed2a40a70997084044aa Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Sun, 12 May 2024 20:26:55 +0800 Subject: [PATCH] dev(narugo): add wd14 syncer, ci skip --- zoo/wd14/sync.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/zoo/wd14/sync.py b/zoo/wd14/sync.py index 8c1785c87e5..4b9ef4a331b 100644 --- a/zoo/wd14/sync.py +++ b/zoo/wd14/sync.py @@ -15,6 +15,7 @@ from tqdm import tqdm from imgutils.tagging.wd14 import MODEL_NAMES +from imgutils.utils import open_onnx_model logging.try_init_root(logging.INFO) @@ -104,7 +105,7 @@ def _is_fc(name): onnx.save_model(model, onnx_file) logging.info(f'Loading and testing for the exported model {onnx_file!r}.') - session = onnxruntime.InferenceSession(onnx_file) + session = open_onnx_model(onnx_file) assert len(session.get_inputs()) == 1 assert len(session.get_outputs()) == 2 assert session.get_outputs()[1].name == embs_outputs[0]