fix ONNXRuntime cuda test bug (open-mmlab#438)

humu789 · May 7, 2022 · 5231e65 · 5231e65
1 parent 86ab063
commit 5231e65
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mmdeploy/codebase/mmdet/deploy/object_detection_model.py b/mmdeploy/codebase/mmdet/deploy/object_detection_model.py
@@ -208,6 +208,7 @@ def forward(self, img: Sequence[torch.Tensor], img_metas: Sequence[dict],
         rescale = kwargs.get('rescale', True)
         for i in range(batch_size):
             dets, labels = batch_dets[i], batch_labels[i]
+            dets = dets.to(device=torch.device(self.device))
             if rescale:
                 scale_factor = img_metas[i]['scale_factor']