code: def get_device_from_env() -> torch.device:
      """Function that gets the torch.device based on the current environment.
  
      This currently supports only CPU, GPU, and MPS devices. If CUDA is available, this function also sets the CUDA device.
  
      Within a distributed context, this function relies on the ``LOCAL_RANK`` environment variable
      to be made available by the program launcher for setting the appropriate device index.
  
      Raises:
          RuntimeError
              If ``LOCAL_RANK`` is outside the range of available GPU devices.
      """
      if torch.cuda.is_available():
          local_rank = int(os.environ.get("LOCAL_RANK", "0"))
          if local_rank >= torch.cuda.device_count():
              raise RuntimeError(
                  "The local rank is larger than the number of available GPUs."
              )
          device = torch.device(f"cuda:{local_rank}")
          torch.cuda.set_device(device)
      elif (
          is_torch_version_geq_1_12()
          and torch.backends.mps.is_built()
          and torch.backends.mps.is_available()
      ):
          device = torch.device("mps")
      else:
          device = torch.device("cpu")
      return device
  
imports: import logging
  import os
  import shutil
  import subprocess
  from collections import defaultdict
  from dataclasses import fields, is_dataclass
  from typing import Any, Dict, Mapping, TypeVar
  import torch
  from torchtnt.utils.version import is_torch_version_geq_1_12
  from typing_extensions import Protocol, runtime_checkable, TypedDict
label_desc: Use image node_id 1 for calling a global function with example usage: get_device_from_env() and returns: device