code: def set_hetero_masks(
      model: torch.nn.Module,
      mask_dict: Dict[EdgeType, Union[Tensor, Parameter]],
      edge_index_dict: Dict[EdgeType, Tensor],
      apply_sigmoid: bool = True,
  ):
      r"""Apply masks to every heterogeneous graph layer in the :obj:`model`
      according to edge types.
      """
      for module in model.modules():
          if isinstance(module, torch.nn.ModuleDict):
              for edge_type in mask_dict.keys():
                  if edge_type in module:
                      edge_level_module = module[edge_type]
                  elif "__".join(edge_type) in module:
                      edge_level_module = module["__".join(edge_type)]
                  else:
                      continue
  
                  set_masks(
                      edge_level_module,
                      mask_dict[edge_type],
                      edge_index_dict[edge_type],
                      apply_sigmoid=apply_sigmoid,
                  )
  
file_path: pytorch_geometric/torch_geometric/explain/algorithm/utils.py
imports: from typing import Dict, Union
  import torch
  from torch import Tensor
  from torch.nn import Parameter
  from torch_geometric.nn import MessagePassing
  from torch_geometric.typing import EdgeType
label_desc: Use image node_id 2 for calling a global function with example usage: set_hetero_masks(model, mask_dict, edge_index_dict, apply_sigmoid) without return types