torchtnt.utils.oom.log_memory_snapshot

torchtnt.utils.oom.log_memory_snapshot(output_dir: str) → None

Writes the memory snapshots to the provided output_dir. For more information, see this blog post .

Parameters:output_dir (str) – The directory to save the memory snapshot.

Note

Outputs are only saved if running on a host with CUDA devices available.