From 58f5dbce7defe0a5421a30c312f6dd81b5db27c9 Mon Sep 17 00:00:00 2001
From: "q.yao" <yaoqian@sensetime.com>
Date: Tue, 10 Aug 2021 20:47:08 +0800
Subject: [PATCH] support cpu deploy_test (#769)

---
 tools/deploy_test.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tools/deploy_test.py b/tools/deploy_test.py
index 51f16b4..56fd61c 100644
--- a/tools/deploy_test.py
+++ b/tools/deploy_test.py
@@ -53,6 +53,7 @@ class ONNXRuntimeSegmentor(BaseSegmentor):
             self.io_binding.bind_output(name)
         self.cfg = cfg
         self.test_mode = cfg.model.test_cfg.mode
+        self.is_cuda_available = is_cuda_available
 
     def extract_feat(self, imgs):
         raise NotImplementedError('This method is not implemented.')
@@ -65,6 +66,10 @@ class ONNXRuntimeSegmentor(BaseSegmentor):
 
     def simple_test(self, img: torch.Tensor, img_meta: Iterable,
                     **kwargs) -> list:
+        if not self.is_cuda_available:
+            img = img.detach().cpu()
+        elif self.device_id >= 0:
+            img = img.cuda(self.device_id)
         device_type = img.device.type
         self.io_binding.bind_input(
             name='input',