code: def convert_captum_output(
      captum_attrs: Tuple[Tensor, ...],
      mask_type: Union[str, MaskLevelType],
      metadata: Optional[Metadata] = None,
  ):
      r"""Convert the output of `Captum.ai <https://captum.ai/>`_ attribution
      methods which is a tuple of attributions to either
      :obj:`(node_mask, edge_mask)` or :obj:`(node_mask_dict, edge_mask_dict)`.
      """
      mask_type = MaskLevelType(mask_type)
      if metadata is not None:
          return captum_output_to_dicts(
              captum_attrs, mask_type, metadata
          )
  
      node_mask = edge_mask = None
      if mask_type == MaskLevelType.edge:
          edge_mask = captum_attrs[0].squeeze(0)
      elif mask_type == MaskLevelType.node:
          node_mask = captum_attrs[0].squeeze(0)
      else:
          node_mask = captum_attrs[0].squeeze(0)
          edge_mask = captum_attrs[1].squeeze(0)
  
      return node_mask, edge_mask
  
imports: from enum import Enum
  from typing import Dict, Optional, Tuple, Union
  import torch
  from torch import Tensor
  from torch_geometric.explain.algorithm.utils import clear_masks, set_hetero_masks, set_masks
  from torch_geometric.explain.config import ModelConfig, ModelMode, ModelReturnType
  from torch_geometric.typing import EdgeType, Metadata, NodeType
label_desc: Use image node_id 4 for calling a global function with example usage: convert_captum_output(captum_attrs, mask_type, metadata) and returns: node_mask,  edge_mask, captum_output_to_dicts