logo
Browse Source

Fix run.py

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 1 year ago
parent
commit
2a0357f7df
  1. 3
      benchmark/run.py

3
benchmark/run.py

@ -1,5 +1,6 @@
import os
import torch
import numpy
import onnxruntime
import towhee
@ -48,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 = torch.rand((1,) + op.config['input_size'])
dummy_input = numpy.random.randn((1,) + op.config['input_size'])
dim = op(dummy_input).shape[0]
print(f'output dim: {dim}')

Loading…
Cancel
Save