Support cuda

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
3 years ago · 7323c4e9d9
1 changed files with 1 additions and 1 deletions
--- 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()