[FIX] Enable YOLOX training on different devices (#7912)

* Enable yolox training on different devices

* Enable yolox resize test on cpu

Author: solyaH

mmdet/models/detectors/yolox.py

@@ -97,7 +97,7 @@ def forward_train(self,
 
         # random resizing
         if (self._progress_in_iter + 1) % self._random_size_interval == 0:
-            self._input_size = self._random_resize()
+            self._input_size = self._random_resize(device=img.device)
         self._progress_in_iter += 1
 
         return losses
@@ -116,8 +116,8 @@ def _preprocess(self, img, gt_bboxes):
                 gt_bbox[..., 1::2] = gt_bbox[..., 1::2] * scale_y
         return img, gt_bboxes
 
-    def _random_resize(self):
-        tensor = torch.LongTensor(2).cuda()
+    def _random_resize(self, device):
+        tensor = torch.LongTensor(2).to(device)
 
         if self.rank == 0:
             size = random.randint(*self._random_size_range)

tests/test_models/test_forward.py

@@ -673,8 +673,6 @@ def test_inference_detector():
     assert len(result) == 2 and len(result[0]) == num_class
 
 
-@pytest.mark.skipif(
-    not torch.cuda.is_available(), reason='requires CUDA support')
 def test_yolox_random_size():
     from mmdet.models import build_detector
     model = _get_detector_cfg('yolox/yolox_tiny_8x8_300e_coco.py')