class_docstr: A callback for memory snapshot collection during training, saving pickle files to the user-specified directory.
  Uses `Memory Snapshots <https://zdevito.github.io/2022/08/16/memory-snapshots.html>`.
  
  Args:
      output_dir: Directory where to save the memory snapshots.
      memory_snapshot_params: Instance of MemorySnapshotParams which will be passed to MemorySnapshotProfiler.
  
  Note: It is recommended to instantiate this callback **as early as possible** in your training/eval/prediction script,
      ideally before model initialization, to make sure all memory allocation is captured.
code: def on_train_step_end(self, state: State, unit: TTrainUnit) -> None:
      self.memory_snapshot_profiler.step()
  
file_path: tnt/torchtnt/framework/callbacks/memory_snapshot.py
imports: import logging
  from typing import Optional
  from torchtnt.framework.callback import Callback
  from torchtnt.framework.state import State
  from torchtnt.framework.unit import TEvalUnit, TPredictUnit, TTrainUnit
  from torchtnt.utils.memory_snapshot_profiler import MemorySnapshotParams, MemorySnapshotProfiler
label_desc: Use image node_id 2 for calling the MemorySnapshot obj's underlying member method code with example usage: obj.on_train_step_end(state, unit) without return types