Fix run.py

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
1 year ago · bd669ec553
1 changed files with 1 additions and 1 deletions
--- 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}')