code: def test_set_clear_mask(hetero_data):
      edge_mask_dict = {
          ("paper", "to", "paper"): torch.ones(200),
          ("author", "to", "paper"): torch.ones(100),
          ("paper", "to", "author"): torch.ones(100),
      }
  
      model = HeteroModel()
  
      set_hetero_masks(
          model, edge_mask_dict, hetero_data.edge_index_dict
      )
      for edge_type in hetero_data.edge_types:
          # Check that masks are correctly set:
          assert torch.allclose(
              model.conv1.convs[edge_type]._edge_mask,
              edge_mask_dict[edge_type],
          )
          assert model.conv1.convs[edge_type].explain
  
      clear_masks(model)
      for edge_type in hetero_data.edge_types:
          assert model.conv1.convs[edge_type]._edge_mask is None
          assert not model.conv1.convs[edge_type].explain
  
      model = to_hetero(
          GraphSAGE(), hetero_data.metadata(), debug=False
      )
  
      set_hetero_masks(
          model, edge_mask_dict, hetero_data.edge_index_dict
      )
      for edge_type in hetero_data.edge_types:
          # Check that masks are correctly set:
          str_edge_type = "__".join(edge_type)
          assert torch.allclose(
              model.conv1[str_edge_type]._edge_mask,
              edge_mask_dict[edge_type],
          )
          assert model.conv1[str_edge_type].explain
  
      clear_masks(model)
      for edge_type in hetero_data.edge_types:
          str_edge_type = "__".join(edge_type)
          assert model.conv1[str_edge_type]._edge_mask is None
          assert not model.conv1[str_edge_type].explain
  
file_path: pytorch_geometric/test/explain/algorithm/test_explain_algorithm_utils.py
imports: import torch
  from torch_geometric.explain.algorithm.utils import clear_masks, set_hetero_masks
  from torch_geometric.nn import GCNConv, HeteroConv, SAGEConv, to_hetero
label_desc: Use image node_id 1 for calling a global function with example usage: test_set_clear_mask(hetero_data) without return types