code: def gpu_memory_usage(device=0):
      return torch.cuda.memory_allocated(device) / 1024.0**3
  
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 2 for calling a global function with example usage: gpu_memory_usage(device) and returns: unknown