Apply suggestions from code review

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

mmengine/dist/dist.py

@@ -790,11 +790,11 @@ def _gather_object(obj: Any,
     is_nccl_backend = group_backend == torch_dist.Backend.NCCL
     is_mccl_backend = group_backend == 'mccl'
     if is_nccl_backend:
-        current_device = torch.device('', torch.cuda.current_device())
+        current_device = torch.device('cuda', torch.cuda.current_device())
         input_tensor = input_tensor.to(current_device)
         local_size = local_size.to(current_device)
     elif is_mccl_backend:
-        current_device = torch.device('', torch.musa.current_device())
+        current_device = torch.device('musa', torch.musa.current_device())
         input_tensor = input_tensor.to(current_device)
         local_size = local_size.to(current_device)
     # Gather all local sizes. This is so that we can find the max size, and

mmengine/dist/utils.py

@@ -11,8 +11,7 @@
 from torch import Tensor
 from torch import distributed as torch_dist
 from torch.distributed import ProcessGroup
-from mmengine.device import is_mlu_available, is_npu_available
-from mmengine.device import is_musa_available
+from mmengine.device import (is_mlu_available, is_npu_available, is_musa_available)
 
 from collections.abc import Iterable, Mapping
 

mmengine/structures/base_data_element.py

@@ -512,7 +512,7 @@ def cuda(self) -> 'BaseDataElement':
 
     # Tensor-like methods
     def musa(self) -> 'BaseDataElement':
-        """Convert all tensors to GPU in data."""
+        """Convert all tensors to musa in data."""
         new_data = self.new()
         for k, v in self.items():
             if isinstance(v, (torch.Tensor, BaseDataElement)):

mmengine/utils/dl_utils/collect_env.py

@@ -10,7 +10,7 @@
 import torch
 
 import mmengine
-from mmengine.device import is_musa_available
+from mmengine.device import is_cuda_available, is_musa_available
 from .parrots_wrapper import TORCH_VERSION, get_build_config, is_rocm_pytorch
 
 
@@ -57,7 +57,7 @@ def collect_env():
     env_info['sys.platform'] = sys.platform
     env_info['Python'] = sys.version.replace('\n', '')
 
-    cuda_available = torch.cuda.is_available()
+    cuda_available = is_cuda_available()
     musa_available = is_musa_available()
     env_info['CUDA available'] = cuda_available
     env_info['MUSA available'] = musa_available

mmengine/utils/dl_utils/time_counter.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from mmengine.device import is_musa_available
+from mmengine.device import is_cuda_available, is_musa_available
 from mmengine.dist.utils import master_only
 from mmengine.logging import MMLogger, print_log
 
@@ -86,7 +86,7 @@ def wrapper(*args, **kwargs):
             self.__count += 1
 
             if self.with_sync:
-                if torch.cuda.is_available():
+                if is_cuda_available():
                     torch.cuda.synchronize()
                 elif is_musa_available():
                     torch.musa.synchronize()
@@ -95,7 +95,7 @@ def wrapper(*args, **kwargs):
             result = fn(*args, **kwargs)
 
             if self.with_sync:
-                if torch.cuda.is_available():
+                if is_cuda_available():
                     torch.cuda.synchronize()
                 elif is_musa_available():
                     torch.musa.synchronize()

tests/test_runner/test_log_processor.py

@@ -7,7 +7,7 @@
 import torch
 from parameterized import parameterized
 
-from mmengine.device import is_musa_available
+from mmengine.device import is_cuda_available, is_musa_available
 from mmengine.logging import HistoryBuffer, MessageHub, MMLogger
 from mmengine.runner import LogProcessor
 from mmengine.testing import RunnerTestCase
@@ -114,7 +114,7 @@ def test_get_log_after_iter(self, by_epoch, mode, log_with_hierarchy):
                         f"time: {train_logs['time']:.4f}  "
                         f"data_time: {train_logs['data_time']:.4f}  ")
 
-            if torch.cuda.is_available() or is_musa_available():
+            if is_cuda_available() or is_musa_available():
                 log_str += 'memory: 100  '
             if mode == 'train':
                 log_str += f"loss_cls: {train_logs['loss_cls']:.4f}"
@@ -142,7 +142,7 @@ def test_get_log_after_iter(self, by_epoch, mode, log_with_hierarchy):
                         f"time: {train_logs['time']:.4f}  "
                         f"data_time: {train_logs['data_time']:.4f}  ")
 
-            if torch.cuda.is_available() or is_musa_available():
+            if is_cuda_available() or is_musa_available():
                 log_str += 'memory: 100  '
 
             if mode == 'train':