From 7323c4e9d9a6009b349144f816ac6de11d521a8f Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Fri, 30 Sep 2022 17:45:34 +0800 Subject: [PATCH] Support cuda Signed-off-by: Jael Gu --- detectron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/detectron.py b/detectron.py index f15d566..cf5b879 100644 --- a/detectron.py +++ b/detectron.py @@ -54,7 +54,7 @@ class Detectron2(NNOperator): # Detectron2 uses BGR-formatted images res = self._predictor(image[:,:,::-1]) res = res['instances'] - boxes = res.get('pred_boxes').tensor.numpy() + boxes = res.get('pred_boxes').tensor.cpu().numpy() classes = res.get('pred_classes').numpy() scores = res.get('scores').numpy()