code: def log_memory_snapshot(
      output_dir: str, file_prefix: Optional[str] = None
  ) -> None:
      """Writes the memory snapshots to the provided ``output_dir``.
      For more information, see this `blog post <https://zdevito.github.io/2022/08/16/memory-snapshots.html>`_ .
  
      Args:
          output_dir: The directory to save the memory snapshot.
          file_prefix: Prefix to append to the filename. If not set, defaults to "memory_snapshot".
  
      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 logging snapshot")
          return
      if not is_torch_version_geq_2_0():
          logger.warning(
              "CUDA memory snapshot utilities are unavailable. Not logging snapshot"
          )
          return
  
      rank = get_global_rank()
      if file_prefix is None:
          file_prefix = "memory_snapshot"
      save_dir = os.path.join(output_dir, f"{file_prefix}_rank{rank}")
      try:
          snapshot = torch.cuda.memory._snapshot()
          _dump_snapshot(save_dir, snapshot)
          logger.info(f"Logged memory snapshot to {save_dir}")
      except Exception as e:
          logger.error(
              f"Failed to log memory snapshot to {save_dir}: {e}"
          )
  
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 6 for calling a global function with example usage: log_memory_snapshot(output_dir, file_prefix) without return types