code: def choose_device(cfg):
      def get_device():
          try:
              if torch.cuda.is_available():
                  return f"cuda:{cfg.local_rank}"
  
              if torch.backends.mps.is_available():
                  return "mps"
  
              raise SystemError("No CUDA/mps device found")
          except Exception:  # pylint: disable=broad-exception-caught
              return "cpu"
  
      cfg.device = get_device()
      if cfg.world_size == 1:
          cfg.device_map = cfg.device_map or "auto"
      else:
          if cfg.device.startswith("cuda"):
              cfg.device_map = {"": torch.cuda.current_device()}
          else:
              cfg.device_map = {"": cfg.device}
  
      # in `accelerate launch`, we need to not pass through any device map and let
      # accelerate figure out which parts of the model to put on which gpu
      accelerate_vars = [
          var for var in os.environ if var.startswith("ACCELERATE_USE_")
      ]
      if accelerate_vars:
          cfg.device_map = None
  
file_path: axolotl/src/axolotl/utils/config.py
imports: import logging
  import os
  import torch
  from transformers.utils import is_torch_bf16_gpu_available
  from axolotl.utils.bench import log_gpu_memory_usage
  from axolotl.utils.models import load_model_config
label_desc: Use image node_id 1 for calling a global function with example usage: choose_device(cfg) and returns: str, cfg, str, str