From 5f283beeb6527d1f61b84a836dbf25d3b402149f Mon Sep 17 00:00:00 2001 From: Hiram Foster Date: Mon, 25 Mar 2024 14:55:48 -0400 Subject: [PATCH] use device arg function has a arg for device which defaults to 'cuda' --- mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py index ebe55e6dbe..627182cb06 100644 --- a/mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py @@ -163,7 +163,7 @@ def _mask_predict_by_feat_single(self, mask_feat, kernels, priors): batch_size = priors.shape[0] hw = mask_feat.size()[-2:] coord = self.prior_generator.single_level_grid_priors( - hw, level_idx=0).to(mask_feat.device) + hw, level_idx=0, device=mask_feat.device) coord = coord.unsqueeze(0).unsqueeze(0).repeat(batch_size, 1, 1, 1) priors = priors.unsqueeze(2) points = priors[..., :2]