From fa1e7179804b9778d1e1b910ec72c58c745a7ef9 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Wed, 14 Dec 2022 19:32:15 +0800 Subject: [PATCH] Fix run.py Signed-off-by: Jael Gu --- benchmark/run.py | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmark/run.py b/benchmark/run.py index 69ed532..4b26e99 100644 --- a/benchmark/run.py +++ b/benchmark/run.py @@ -129,6 +129,7 @@ elif args.format == 'onnx': img = op.tfms(img).unsqueeze(0).cpu().detach().numpy() features = sess.run(output_names=['output_0'], input_feed={'input_0': img})[0] if len(features.shape) == 4: + features = torch.from_numpy(features) global_pool = nn.AdaptiveAvgPool2d(1) features = global_pool(features) outs = features.squeeze(0)