code: def gpu_mem_get_mbs(torch_gpu_id=None):
      """Query nvidia for total, used and free memory for gpu in MBs. if gpu id is not passed, currently selected torch device is used"""
      if not use_gpu:
          return GPUMemory(0, 0, 0)
      if torch_gpu_id is None:
          torch_gpu_id = torch.cuda.current_device()
      nvml_gpu_id = get_nvml_gpu_id(torch_gpu_id)
      try:
          handle = pynvml.nvmlDeviceGetHandleByIndex(nvml_gpu_id)
          info = pynvml.nvmlDeviceGetMemoryInfo(handle)
          return GPUMemory(
              *(map(b2mb, [info.total, info.free, info.used]))
          )
      except:
          return GPUMemory(0, 0, 0)
  
file_path: ipyexperiments/ipyexperiments/utils/mem.py
imports: import gc
  import os
  from collections import namedtuple
  from ipyexperiments.utils.pynvml_gate import load_pynvml_env
label_desc: Use image node_id 4 for calling a global function with example usage: gpu_mem_get_mbs(torch_gpu_id) and returns: GPUMemory, GPUMemory, GPUMemory