From 621ee3b17c7f8efe5d6585c3a08799fb9ddc5f9a Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Fri, 30 Sep 2022 17:46:51 +0800 Subject: [PATCH] Support cuda Signed-off-by: Jael Gu --- detectron.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/detectron.py b/detectron.py index cf5b879..19fb84f 100644 --- a/detectron.py +++ b/detectron.py @@ -55,7 +55,7 @@ class Detectron2(NNOperator): res = self._predictor(image[:,:,::-1]) res = res['instances'] boxes = res.get('pred_boxes').tensor.cpu().numpy() - classes = res.get('pred_classes').numpy() - scores = res.get('scores').numpy() + classes = res.get('pred_classes').cpu().numpy() + scores = res.get('scores').cpu().numpy() return (boxes, classes, scores)