code: def GetGPUMemoryUsageStats():
      """Get GPU memory usage stats from CUDAContext/HIPContext. This requires flag
      --caffe2_gpu_memory_tracking to be enabled"""
      from caffe2.python import workspace, core
  
      workspace.RunOperatorOnce(
          core.CreateOperator(
              "GetGPUMemoryUsage",
              [],
              ["____mem____"],
              device_option=core.DeviceOption(
                  workspace.GpuDeviceType, 0
              ),
          ),
      )
      b = workspace.FetchBlob("____mem____")
      return {
          "total_by_gpu": b[0, :],
          "max_by_gpu": b[1, :],
          "total": np.sum(b[0, :]),
          "max_total": np.sum(b[1, :]),
      }
  
imports: from caffe2.proto import caffe2_pb2
  from google.protobuf.message import DecodeError, Message
  from google.protobuf import text_format
  import sys
  import collections
  import copy
  import functools
  import numpy
label_desc: Use image node_id 10 for calling a global function with example usage: GetGPUMemoryUsageStats() and returns: dict