From bd669ec55387ae77072515bf6f624731d5f8d3a7 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Wed, 14 Dec 2022 19:09:08 +0800 Subject: [PATCH] Fix run.py Signed-off-by: Jael Gu --- benchmark/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/run.py b/benchmark/run.py index 6857868..b28f436 100644 --- a/benchmark/run.py +++ b/benchmark/run.py @@ -49,7 +49,7 @@ query_data = query_data[:query_size] if query_size else query_data # Warm up print('Warming up...') op = ops.image_embedding.timm(model_name=model_name, device=device).get_op() -dummy_input = numpy.random.randn((1,) + op.config['input_size']) +dummy_input = numpy.random.randn(1, 3, op.config['input_size'][-2], op.config['input_size'][-1]) dim = op(dummy_input).shape[0] print(f'output dim: {dim}')