sparsity = 0.7
import matplotlib.pyplot as plt

# Visualization
gnnexplainer_related_preds = \
    gnnexplainer(data.x, data.edge_index, sparsity=sparsity, num_classes=num_classes, node_idx=node_idx)
ax, G = gnnexplainer.visualize_graph(node_idx=node_idx,
                                     edge_index=data.edge_index,
                                     edge_mask=gnnexplainer_related_preds[1][prediction],
                                     y=data.y)
plt.show()


TypeError                                 Traceback (most recent call last)
[<ipython-input-8-413a9bd50c1c>](https://localhost:8080/#) in <cell line: 7>()
      5 gnnexplainer_related_preds = \
      6     gnnexplainer(data.x, data.edge_index, sparsity=sparsity, num_classes=num_classes, node_idx=node_idx)
----> 7 ax, G = gnnexplainer.visualize_graph(node_idx=node_idx,
      8                                      edge_index=data.edge_index,
      9                                      edge_mask=gnnexplainer_related_preds[1][prediction],

[/usr/local/lib/python3.10/dist-packages/dig/xgraph/method/base_explainer.py](https://localhost:8080/#) in visualize_graph(self, node_idx, edge_index, edge_mask, y, threshold, nolabel, **kwargs)
    240                     connectionstyle="arc3,rad=0.08",  # rad control angle
    241                 ))
--> 242         nx.draw_networkx_nodes(G, pos, node_color=node_colors, **kwargs)
    243         # define node labels
    244         if self.molecule:

TypeError: draw_networkx_nodes() got an unexpected keyword argument 'with_labels'