code: def get_gpu_memory_usage(ctx: List[mx.context.Context]) -> Dict[int, Tuple[int, int]]:
      
      if isinstance(ctx, mx.context.Context):
          ctx = [ctx]
      ctx = [c for c in ctx if c.device_type == 'gpu']
      if not ctx:
          return {}
      if shutil.which("nvidia-smi") is None:
          logger.warning("Couldn't find nvidia-smi, therefore we assume no GPUs are available.")
          return {}
  
      device_ids = [c.device_id for c in ctx]
  
      # Run from clean forkserver process to not leak any CUDA resources
  
      mp_context = mp_utils.get_context()
      result_queue = mp_context.Queue()
      nvidia_smi_process = mp_context.Process(target=query_nvidia_smi, args=(device_ids, result_queue,))
      nvidia_smi_process.start()
      nvidia_smi_process.join()
  
      memory_data = result_queue.get()
  
      log_gpu_memory_usage(memory_data)
  
      return memory_data
docs: Returns used and total memory for GPUs identified by the given context list.
  
      :param ctx: List of MXNet context devices.
      :return: Dictionary of device id mapping to a tuple of (memory used, memory total).