code: def forward(
      self, x: Tensor, edge_index: Tensor, **kwargs
  ) -> Tuple[List, List, List[Dict]]:
      # set default subgraph with 10 edges
  
      pred_label = kwargs.get("pred_label")
      num_classes = kwargs.get("num_classes")
      self.model.eval()
      self.explainer.__clear_masks__()
  
      x = x.to(self.device)
      edge_index, _ = add_remaining_self_loops(
          edge_index, num_nodes=x.size(0)
      )
      edge_index = edge_index.to(self.device)
  
      if self.explain_graph:
          embed = self.model.get_emb(x, edge_index)
          _, edge_mask = self.explainer.explain(
              x, edge_index, embed=embed, tmp=1.0, training=False
          )
          # edge_masks
          edge_masks = [edge_mask for _ in range(num_classes)]
          # Calculate mask
          hard_edge_masks = [
              self.control_sparsity(
                  edge_mask, sparsity=kwargs.get("sparsity")
              ).sigmoid()
              for _ in range(num_classes)
          ]
  
          self.__clear_masks__()
          self.__set_masks__(x, edge_index)
          with torch.no_grad():
              if self.explain_graph:
                  related_preds = self.eval_related_pred(
                      x, edge_index, hard_edge_masks
                  )
  
          self.__clear_masks__()
  
      else:
          node_idx = kwargs.get("node_idx")
          sparsity = kwargs.get("sparsity")
          assert (
              kwargs.get("node_idx") is not None
          ), "please input the node_idx"
          select_edge_index = torch.arange(0, edge_index.shape[1])
          (
              subgraph_x,
              subgraph_edge_index,
              _,
              subset,
              kwargs,
          ) = self.explainer.get_subgraph(
              node_idx,
              x,
              edge_index,
              select_edge_index=select_edge_index,
          )
          select_edge_index = kwargs["select_edge_index"]
          self.select_edge_mask = edge_index.new_empty(
              edge_index.size(1), device=self.device, dtype=torch.bool
          )
          self.select_edge_mask.fill_(False)
          self.select_edge_mask[select_edge_index] = True
          self.hard_edge_mask = edge_index.new_empty(
              subgraph_edge_index.size(1),
              device=self.device,
              dtype=torch.bool,
          )
          self.hard_edge_mask.fill_(True)
          self.subset = subset
          self.new_node_idx = torch.where(subset == node_idx)[0]
  
          subgraph_embed = self.model.get_emb(
              subgraph_x, subgraph_edge_index
          )
          _, subgraph_edge_mask = self.explainer.explain(
              subgraph_x,
              subgraph_edge_index,
              embed=subgraph_embed,
              tmp=1.0,
              training=False,
              node_idx=self.new_node_idx,
          )
  
          # edge_masks
          edge_masks = [subgraph_edge_mask for _ in range(num_classes)]
          # Calculate mask
          hard_edge_masks = [
              self.control_sparsity(
                  subgraph_edge_mask, sparsity=sparsity
              ).sigmoid()
              for _ in range(num_classes)
          ]
  
          self.__clear_masks__()
          self.__set_masks__(subgraph_x, subgraph_edge_index)
          with torch.no_grad():
              related_preds = self.eval_related_pred(
                  subgraph_x,
                  subgraph_edge_index,
                  hard_edge_masks,
                  node_idx=self.new_node_idx,
              )
  
          self.__clear_masks__()
  
      return edge_masks, hard_edge_masks, related_preds
  
imports: import os
  import torch
  import hydra
  from omegaconf import OmegaConf
  from sklearn.metrics import roc_auc_score
  from dig.xgraph.method import PGExplainer
  from dig.xgraph.evaluation import XCollector
  from dig.xgraph.method.base_explainer import ExplainerBase
  from dig.xgraph.utils.compatibility import compatible_state_dict
  from benchmarks.xgraph.gnnNets import get_gnnNets
  from benchmarks.xgraph.dataset import get_dataset, get_dataloader, SynGraphDataset
  from benchmarks.xgraph.utils import check_dir, fix_random_seed, Recorder, perturb_input
  from torch import Tensor
  from typing import List, Dict, Tuple
  from torch_geometric.utils import add_remaining_self_loops
label_desc: Use image node_id 2 for calling the PGExplainer_edges obj's underlying member method code with example usage: obj.forward(x, edge_index) and returns: edge_masks,  hard_edge_masks,  related_preds