code: def __init__(self, *, use_cuda: bool, cuda_device: int):
      self.use_cuda = use_cuda
      self.cuda_device = cuda_device
      self.cuda_count = torch.cuda.device_count()
  
      self.is_cuda = self.use_cuda and torch.cuda.is_available()
      if not self.is_cuda:
          self.device = torch.device("cpu")
      else:
          if self.cuda_device < self.cuda_count:
              self.device = torch.device("cuda", self.cuda_device)
          else:
              self.device = torch.device("cuda", self.cuda_count - 1)
  
file_path: labml/helpers/labml_helpers/device.py
imports: import torch
  from labml.configs import BaseConfigs, hyperparams, option
label_desc: Use image node_id 1 to create a new DeviceInfo object  with example: obj = DeviceInfo()