[Refactor] Handle case where device is neither CPU nor CUDA in HamHead (#2868)

Author: KKIEEK

mmseg/models/decode_heads/ham_head.py

@@ -5,6 +5,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from mmcv.cnn import ConvModule
+from mmengine.device import get_device
 
 from mmseg.registry import MODELS
 from ..utils import resize
@@ -52,7 +53,7 @@ def __init__(self,
 
         self.rand_init = rand_init
 
-    def _build_bases(self, B, S, D, R, cuda=False):
+    def _build_bases(self, B, S, D, R, device=None):
         raise NotImplementedError
 
     def local_step(self, x, bases, coef):
@@ -80,14 +81,13 @@ def forward(self, x, return_bases=False):
         D = C // self.S
         N = H * W
         x = x.view(B * self.S, D, N)
-        cuda = 'cuda' in str(x.device)
         if not self.rand_init and not hasattr(self, 'bases'):
-            bases = self._build_bases(1, self.S, D, self.R, cuda=cuda)
+            bases = self._build_bases(1, self.S, D, self.R, device=x.device)
             self.register_buffer('bases', bases)
 
         # (S, D, R) -> (B * S, D, R)
         if self.rand_init:
-            bases = self._build_bases(B, self.S, D, self.R, cuda=cuda)
+            bases = self._build_bases(B, self.S, D, self.R, device=x.device)
         else:
             bases = self.bases.repeat(B, 1, 1)
 
@@ -116,13 +116,11 @@ def __init__(self, args=dict()):
 
         self.inv_t = 1
 
-    def _build_bases(self, B, S, D, R, cuda=False):
+    def _build_bases(self, B, S, D, R, device=None):
         """Build bases in initialization."""
-        if cuda:
-            bases = torch.rand((B * S, D, R)).cuda()
-        else:
-            bases = torch.rand((B * S, D, R))
-
+        if device is None:
+            device = get_device()
+        bases = torch.rand((B * S, D, R)).to(device)
         bases = F.normalize(bases, dim=1)
 
         return bases