code: def test_log_memory_snapshot(self) -> None:
      with tempfile.TemporaryDirectory() as temp_dir:
          # Record history
          torch.cuda.memory._record_memory_history(
              enabled="all", max_entries=10000
          )
  
          # initialize device & allocate memory for tensors
          device = get_device_from_env()
          a = torch.rand((1024, 1024), device=device)
          b = torch.rand((1024, 1024), device=device)
          _ = (a + b) * (a - b)
  
          # save a snapshot
          log_memory_snapshot(temp_dir, "foo")
  
          # Check if the corresponding files exist
          save_dir = os.path.join(temp_dir, "foo_rank0")
  
          pickle_dump_path = os.path.join(save_dir, "snapshot.pickle")
          self.assertTrue(os.path.exists(pickle_dump_path))
  
          trace_path = os.path.join(save_dir, "trace_plot.html")
          self.assertTrue(os.path.exists(trace_path))
  
          segment_plot_path = os.path.join(
              save_dir, "segment_plot.html"
          )
          self.assertTrue(os.path.exists(segment_plot_path))
  
file_path: tnt/tests/utils/test_oom.py
imports: import os
  import tempfile
  import unittest
  import torch
  from torchtnt.utils.device import get_device_from_env
  from torchtnt.utils.oom import _bytes_to_mb_gb, is_out_of_cpu_memory, is_out_of_cuda_memory, is_out_of_memory_error, log_memory_snapshot
  from torchtnt.utils.version import is_torch_version_geq_2_0
label_desc: Use image node_id 4 for calling the OomTest obj's underlying member method code with example usage: obj.test_log_memory_snapshot() without return types