GNNExplainer explain_graph function throw error in sample code

[ChengyuanSha]

Hello community,

As I'm testing explain_graph function in GNNExplainer using the sample code provided, an error occurred. I'm not sure what's the cause. Could you please help me solve this bug?

Thanks in advance!!

import os.path as osp

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.nn import GCNConv, GNNExplainer

dataset = 'Cora'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Planetoid')
transform = T.Compose([T.GCNNorm(), T.NormalizeFeatures()])
dataset = Planetoid(path, dataset, transform=transform)
data = dataset[0]


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16, normalize=False)
        self.conv2 = GCNConv(16, dataset.num_classes, normalize=False)

    def forward(self, x, edge_index, edge_weight):
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index, edge_weight)
        return F.log_softmax(x, dim=1)


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight

for epoch in range(1, 201):
    model.train()
    optimizer.zero_grad()
    log_logits = model(x, edge_index, edge_weight)
    loss = F.nll_loss(log_logits[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

explainer = GNNExplainer(model, epochs=200, return_type='log_prob')
explainer.explain_graph(x, edge_index)

It threw

Traceback (most recent call last):
  File "----explainer_test.py", line 45, in <module>
    explainer.explain_graph(x, edge_index)
  File "---\anaconda3\envs\AI-CUDA\lib\site-packages\torch_geometric\nn\models\gnn_explainer.py", line 205, in explain_graph
    out = self.model(x=x, edge_index=edge_index, batch=batch, **kwargs)
  File "----\anaconda3\envs\AI-CUDA\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'batch'

[wsad1]

The Planetoid(dataset='Cora') is a node level classification task so you need to use the explain_node function to explain a nodes prediction.
explain_graph is used to explain graph level tasks (classifications, regression) .