code: def set_masks(
      model: torch.nn.Module,
      mask: Union[Tensor, Parameter],
      edge_index: Tensor,
      apply_sigmoid: bool = True,
  ):
      r"""Apply mask to every graph layer in the :obj:`model`."""
      loop_mask = edge_index[0] != edge_index[1]
  
      # Loop over layers and set masks on MessagePassing layers:
      for module in model.modules():
          if isinstance(module, MessagePassing):
              # Skip layers that have been explicitly set to `False`:
              if module.explain is False:
                  continue
  
              # Convert mask to a param if it was previously registered as one.
              # This is a workaround for the fact that PyTorch does not allow
              # assignments of pure tensors to parameter attributes:
              if (
                  not isinstance(mask, Parameter)
                  and "_edge_mask" in module._parameters
              ):
                  mask = Parameter(mask)
  
              module.explain = True
              module._edge_mask = mask
              module._loop_mask = loop_mask
              module._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 1 for calling a global function with example usage: set_masks(model, mask, edge_index, apply_sigmoid) without return types