From 91558791340fb21b38064d6ba772936015c2fc31 Mon Sep 17 00:00:00 2001 From: maningsheng Date: Fri, 5 Mar 2021 17:45:36 +0800 Subject: [PATCH] fix test unit of nms and batched_nms for tensorrt --- tests/test_ops/test_tensorrt.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_ops/test_tensorrt.py b/tests/test_ops/test_tensorrt.py index 577b43d86b..09f706f206 100644 --- a/tests/test_ops/test_tensorrt.py +++ b/tests/test_ops/test_tensorrt.py @@ -123,8 +123,8 @@ def test_nms(): fp16_mode = False max_workspace_size = 1 << 30 data = mmcv.load('./tests/data/batched_nms_data.pkl') - boxes = data['boxes'].cuda() - scores = data['scores'].cuda() + boxes = torch.from_numpy(data['boxes']).cuda() + scores = torch.from_numpy(data['scores']).cuda() nms = partial(nms, iou_threshold=0.7, offset=0) wrapped_model = WrapFunction(nms) wrapped_model.cpu().eval() @@ -195,9 +195,9 @@ def test_batched_nms(): max_workspace_size = 1 << 30 data = mmcv.load('./tests/data/batched_nms_data.pkl') nms_cfg = dict(type='nms', iou_threshold=0.7) - boxes = data['boxes'].cuda() - scores = data['scores'].cuda() - idxs = data['idxs'].cuda() + boxes = torch.from_numpy(data['boxes']).cuda() + scores = torch.from_numpy(data['scores']).cuda() + idxs = torch.from_numpy(data['idxs']).cuda() class_agnostic = False nms = partial(batched_nms, nms_cfg=nms_cfg, class_agnostic=class_agnostic)