code: def sample(self):
      cuda_mem = StatefulTensor.GST_MGR.total_mem["cuda"]
      cpu_mem = StatefulTensor.GST_MGR.total_mem["cpu"]
      timestamp = get_timestamp_us()
      if self._tracing:
          self.events.append(
              StatefulTensorMemoryEvent(
                  timestamp, DeviceType.CUDA, cuda_mem
              )
          )
          self.events.append(
              StatefulTensorMemoryEvent(
                  timestamp, DeviceType.CPU, cpu_mem
              )
          )
  
imports: import os
  import threading
  import time
  from enum import Enum
  from typing import List
  import torch
  from colossalai.gemini.ophooks import BaseOpHook
  from colossalai.gemini.stateful_tensor import StatefulTensor
  from colossalai.legacy.engine import Engine
  from colossalai.legacy.utils.profiler.extention import ProfilerExtension
label_desc: Use image node_id 2 for calling the StatefulTensorMemoryTracer obj's underlying member method code with example usage: obj.sample() without return types