class_docstr: Records a history of memory allocation and free events, and dumps to a
  file which can be visualized offline. It by default keeps track of
  100000 events, and dumps at user specified step as well as on OOM.
  The profiler is ideally started before model is instantiated in memory
  so that allocations that stay constant in training are accounted for.
  
  Args:
      output_dir: Directory where to save the memory snapshots.
      memory_snapshot_params: Instance of MemorySnapshotParams.
  
  Raises:
      ValueError: If `start_step` is negative, or `stop_step` is less than or equal to zero.
      ValueError: If `start_step` is greater than or equal to `stop_step`.
      ValueError: If `start_step` is set and `stop_step` is not set.
      ValueError: If `stop_step` is set and neither `start_step` nor `enable_oom_observer` are set.
      ValueError: If `enable_oom_observer` is False and neither `start_step` nor `stop_step` is set
  
  Examples::
      memory_snapshot_params = MemorySnapshotParams(start_step=5, stop_step=10, enable_oom_observer=True)
      memory_snapshot_profiler = MemorySnapshotProfiler(output_dir="/tmp", memory_snapshot_params=memory_snapshot_params)
      for batch in dataloader:
          ...
          memory_snapshot_profiler.step()
class_docstr_tok: ["Records","a","history","of","memory","allocation","and","free","events",",","and","dumps","to","a","file","which","can","be","visualized","offline",".","It","by","default","keeps","track","of","100000","events",",","and","dumps","at","user","specified","step","as","well","as","on","OOM",".","The","profiler","is","ideally","started","before","model","is","instantiated","in","memory","so","that","allocations","that","stay","constant","in","training","are","accounted","for",".","Args",":","output_dir",":","Directory","where","to","save","the","memory","snapshots",".","memory_snapshot_params",":","Instance","of","MemorySnapshotParams",".","Raises",":","ValueError",":","If","`","start_step","`","is","negative",",","or","`","stop_step","`","is","less","than","or","equal","to","zero",".","ValueError",":","If","`","start_step","`","is","greater","than","or","equal","to","`","stop_step","`",".","ValueError",":","If","`","start_step","`","is","set","and","`","stop_step","`","is","not","set",".","ValueError",":","If","`","stop_step","`","is","set","and","neither","`","start_step","`","nor","`","enable_oom_observer","`","are","set",".","ValueError",":","If","`","enable_oom_observer","`","is","False","and","neither","`","start_step","`","nor","`","stop_step","`","is","set","Examples",":",":","memory_snapshot_params","=","MemorySnapshotParams","(","start_step=5",",","stop_step=10",",","enable_oom_observer=True",")","memory_snapshot_profiler","=","MemorySnapshotProfiler","(","output_dir=","''","\/tmp","''",",","memory_snapshot_params=memory_snapshot_params",")","for","batch","in","dataloader",":","...","memory_snapshot_profiler.step","(",")"]
code: def __exit__(
      self,
      exc_type: Optional[Type[BaseException]],
      exc_value: Optional[BaseException],
      tb: Optional[TracebackType],
  ) -> Optional[bool]:
      self.stop()
  
file_path: tnt/torchtnt/utils/memory_snapshot_profiler.py
imports: import logging
  from dataclasses import dataclass
  from types import TracebackType
  from typing import Optional, Type
  import torch
  from torchtnt.utils.oom import attach_oom_observer, log_memory_snapshot
  from torchtnt.utils.version import is_torch_version_geq_2_0
label_desc: Use image node_id 3 for calling the MemorySnapshotProfiler obj's underlying member method code with example usage: obj.__exit__(exc_type, exc_value, tb) without return types