code: def test_custom_explain_message():
      x = torch.randn(4, 8)
      edge_index = torch.tensor(
          [[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]]
      )
  
      conv = SAGEConv(8, 32)
  
      def explain_message(self, inputs, x_i, x_j):
          assert isinstance(self, SAGEConv)
          assert inputs.size() == (6, 8)
          assert inputs.size() == x_i.size() == x_j.size()
          assert torch.allclose(inputs, x_j)
          self.x_i = x_i
          self.x_j = x_j
          return inputs
  
      conv.explain_message = explain_message.__get__(
          conv, MessagePassing
      )
      conv.explain = True
  
      conv(x, edge_index)
  
      assert torch.allclose(conv.x_i, x[edge_index[1]])
      assert torch.allclose(conv.x_j, x[edge_index[0]])
  
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 3 for calling a global function with example usage: test_custom_explain_message() and returns: inputs