code: def get_device_memory_usage(
      kineto_event: _KinetoEvent,
      memory_events: List[List[Union[_KinetoEvent, bool]]],
  ) -> Dict[str, DeviceUsageMetrics]:
      """Get CPU and CUDA memory usage for an event.
  
      :param kineto_event: a Kineto event instance.
      :param memory_events: list of memory events.
      """
      mem_records_acc = profiler_util.MemRecordsAcc(memory_events)
      records_in_interval = mem_records_acc.in_interval(
          kineto_event.start_us(),
          kineto_event.start_us() + kineto_event.duration_us(),
      )
      memory_so_far = defaultdict(int)
      count_so_far = defaultdict(int)
      average_so_far = defaultdict(float)
      max_so_far = defaultdict(int)
  
      for mem_record in records_in_interval:
          device, nbytes = get_memory_details(mem_record[0])
          memory_so_far[device] += nbytes
          max_so_far[device] = max(
              max_so_far[device], memory_so_far[device]
          )
          average_so_far[device] = (
              memory_so_far[device]
              + (average_so_far[device] * count_so_far[device])
          ) / (count_so_far[device] + 1)
          count_so_far[device] += 1
      memory_info_per_device = {}
      for device in count_so_far:
          memory_info_per_device[
              f"torch_{device}_"
          ] = DeviceUsageMetrics(
              max_memory_used=max_so_far[device],
              average_memory_used=average_so_far[device],
          )
      return memory_info_per_device
  
file_path: ludwig/ludwig/benchmarking/reporting.py
imports: from collections import Counter, defaultdict
  from statistics import mean
  from typing import Any, Dict, List, Tuple, Union
  import torch
  from torch._C._autograd import _KinetoEvent
  from torch.autograd import DeviceType, profiler_util
  from ludwig.benchmarking.profiler_dataclasses import DeviceUsageMetrics, SystemResourceMetrics, TorchProfilerMetrics
  from ludwig.constants import LUDWIG_TAG
label_desc: Use image node_id 3 for calling a global function with example usage: get_device_memory_usage(kineto_event, memory_events) and returns: memory_info_per_device