code: def gpu_memory_usage_smi(device=0):
      if isinstance(device, torch.device):
          device = device.index
      if isinstance(device, str) and device.startswith("cuda:"):
          device = int(device[5:])
      try:
          pynvml.nvmlInit()
          handle = pynvml.nvmlDeviceGetHandleByIndex(device)
          info = pynvml.nvmlDeviceGetMemoryInfo(handle)
          return info.used / 1024.0**3
      except NVMLError:
          return 0.0
  
file_path: axolotl/src/axolotl/utils/bench.py
imports: import functools
  import pynvml
  import torch
  from pynvml.nvml import NVMLError
label_desc: Use image node_id 4 for calling a global function with example usage: gpu_memory_usage_smi(device) and returns: unknown, int