code: def attach_oom_observer(
      output_dir: str, trace_max_entries: int = 1000000
  ) -> None:
      """Attaches a function to record the PyTorch memory snapshot when an out of memory error occurs.
  
      For more information, see this `blog post <https://zdevito.github.io/2022/08/16/memory-snapshots.html>`_ .
  
      Args:
          output_dir (str): The directory to save the memory snapshot.
          trace_max_entries (int, optional): The maximum number of trace entries to record. Defaults to 1000000.
  
      Note:
          Outputs are only saved if running on a host with CUDA devices available.
      """
      if not torch.cuda.is_available():
          logger.info("CUDA unavailable. Not attaching OOM observer.")
          return
      if not is_torch_version_geq_2_0():
          logger.warning(
              "CUDA memory snapshot utilities are unavailable. Not attaching OOM observer."
          )
          return
  
      torch.cuda.memory._record_memory_history(
          enabled="all", max_entries=trace_max_entries
      )
      # pyre-fixme[16]: Module `_C` has no attribute
      #  `_cuda_attach_out_of_memory_observer`.
      torch._C._cuda_attach_out_of_memory_observer(
          _oom_observer(output_dir)
      )
  
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 8 for calling a global function with example usage: attach_oom_observer(output_dir, trace_max_entries) without return types