code: def torch_mem_stats():
      objects = gc.get_objects()
      tensors = [
          obj for obj in objects if torch.is_tensor(obj) and obj.is_cuda
      ]
  
      total_numel = 0
      total_mem = 0
      visited_data = set()
      for tensor in tensors:
          # a data_ptr indicates a memory block allocated
          data_ptr = tensor.storage().data_ptr()
          if data_ptr in visited_data:
              continue
          visited_data.add(data_ptr)
  
          print(tensor.shape, tensor.data_ptr())
  
          numel = tensor.numel()
          total_numel += numel
          element_size = tensor.storage().element_size()
          mem = numel * element_size
          total_mem += mem
  
      return total_mem
  
imports: import argparse
  import dataclasses
  from attr import define, field
  from attr.setters import frozen
  import functools
  import gc
  import math
  import os
  from typing import Tuple, Union, Optional, Any, Sequence, List
  import numpy
  import psutil
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: torch_mem_stats() and returns: total_mem