From 8c9dafe2b49e764df847bfa4706e125ab8ca719a Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Wed, 22 Dec 2021 19:03:53 +0800 Subject: [PATCH] Update the test script Signed-off-by: shiyu22 --- test_image_embedding_pipeline_yaml.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test_image_embedding_pipeline_yaml.py b/test_image_embedding_pipeline_yaml.py index 7bf7cd7..44e8189 100644 --- a/test_image_embedding_pipeline_yaml.py +++ b/test_image_embedding_pipeline_yaml.py @@ -3,5 +3,6 @@ from towhee import pipeline def test_image_embedding_resnet50(repo_name='towhee/image-embedding-pipeline-template', img_path='./test_data/test.jpg', dimension=2048): embedding_pipeline = pipeline(repo_name) - embedding = embedding_pipeline(img_path) - assert (1, dimension)==embedding[0].shape + (embedding,) = embedding_pipeline(img_path) + # get the output shape in README + assert (dimension,)==embedding[0].shape