code: def _oom_observer(
      output_dir: str,
  ) -> Callable[[Union[int, torch.device], int, int, int], None]:
      def oom_logger(
          device: Union[int, torch.device],
          alloc: int,
          device_alloc: int,
          device_free: int,
      ) -> None:
          """
          Log memory snapshot in the event of CUDA OOM.
          """
          logger.info(
              f"Saving memory snapshot device: {device}, alloc: {_bytes_to_mb_gb(alloc)}, device_alloc: {_bytes_to_mb_gb(device_alloc)}, device_free: {_bytes_to_mb_gb(device_free)}"
          )
          try:
              log_memory_snapshot(output_dir, "oom")
          except Exception as e:
              logger.error(
                  f"Failed to log memory snapshot during OOM {e}"
              )
  
      return oom_logger
  
file_path: tnt/torchtnt/utils/oom.py
imports: import logging
  import os
  import pickle
  from typing import Any, Callable, Dict, Optional, Union
  import torch
  from torchtnt.utils.distributed import get_global_rank
  from torchtnt.utils.fsspec import get_filesystem
  from torchtnt.utils.version import is_torch_version_geq_2_0
label_desc: Use image node_id 5 for calling a global function with example usage: _oom_observer(output_dir) and returns: oom_logger