From 36653e1a14fc837c7ff09171d01e324de0043760 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Wed, 14 Dec 2022 15:59:41 +0800 Subject: [PATCH] Update onnx test Signed-off-by: Jael Gu --- test_onnx.py | 5 ++--- test_onnx2.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/test_onnx.py b/test_onnx.py index f08121c..cd95354 100644 --- a/test_onnx.py +++ b/test_onnx.py @@ -1,5 +1,4 @@ -from auto_transformers import AutoTransformers - +from towhee import ops import numpy import onnx import onnxruntime @@ -48,7 +47,7 @@ for name in models: f.write(','.join(status) + '\n') status = [name] + ['fail'] * 5 try: - op = AutoTransformers(model_name=name) + op = ops.text_embedding.transformers(model_name=name, device='cpu').get_op() out1 = op(test_txt) logger.info('OP LOADED.') status[1] = 'success' diff --git a/test_onnx2.py b/test_onnx2.py index d8a981e..47f8740 100644 --- a/test_onnx2.py +++ b/test_onnx2.py @@ -1,4 +1,4 @@ -from auto_transformers import AutoTransformers +from towhee import ops import torch import numpy import onnx @@ -82,7 +82,7 @@ for name in models: f.write(','.join(status) + '\n') status = [name] + ['fail'] * 5 try: - op = AutoTransformers(model_name=name, device='cpu') + op = ops.text_embedding.transformers(model_name=name, device='cpu').get_op() out1 = op(test_txt) logger.info('OP LOADED.') status[1] = 'success'