code: def test_to_captum(model, mask_type, output_idx):
      captum_model = to_captum_model(
          model, mask_type=mask_type, output_idx=output_idx
      )
      pre_out = model(x, edge_index)
      if mask_type == "node":
          mask = x * 0.0
          out = captum_model(mask.unsqueeze(0), edge_index)
      elif mask_type == "edge":
          mask = (
              torch.ones(
                  edge_index.shape[1],
                  dtype=torch.float,
                  requires_grad=True,
              )
              * 0.5
          )
          out = captum_model(mask.unsqueeze(0), x, edge_index)
      elif mask_type == "node_and_edge":
          node_mask = x * 0.0
          edge_mask = (
              torch.ones(
                  edge_index.shape[1],
                  dtype=torch.float,
                  requires_grad=True,
              )
              * 0.5
          )
          out = captum_model(
              node_mask.unsqueeze(0), edge_mask.unsqueeze(0), edge_index
          )
  
      if output_idx is not None:
          assert out.shape == (1, 7)
          assert torch.any(out != pre_out[[output_idx]])
      else:
          assert out.shape == (8, 7)
          assert torch.any(out != pre_out)
  
file_path: pytorch_geometric/test/explain/algorithm/test_captum.py
imports: import pytest
  import torch
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.explain.algorithm.captum import to_captum_input
  from torch_geometric.nn import GAT, GCN, SAGEConv
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.nn.models import to_captum_model
  from torch_geometric.testing import withPackage
label_desc: Use image node_id 1 for calling a global function with example usage: test_to_captum(model, mask_type, output_idx) without return types