code: def extract_step(
      self,
      x: Tensor,
      edge_index: Tensor,
      detach: bool = True,
      split_fc: bool = False,
  ):
      layer_extractor = []
      hooks = []
  
      def register_hook(module: nn.Module):
          if not list(module.children()) or isinstance(
              module, MessagePassing
          ):
              hooks.append(module.register_forward_hook(forward_hook))
  
      def forward_hook(
          module: nn.Module, input: Tuple[Tensor], output: Tensor
      ):
          # input contains x and edge_index
          if detach:
              layer_extractor.append(
                  (
                      module,
                      input[0].clone().detach(),
                      output.clone().detach(),
                  )
              )
          else:
              layer_extractor.append((module, input[0], output))
  
      # --- register hooks ---
      self.model.apply(register_hook)
  
      pred = self.model(x, edge_index)
  
      for hook in hooks:
          hook.remove()
  
      # --- divide layer sets ---
  
      walk_steps = []
      fc_steps = []
      pool_flag = False
      step = {"input": None, "module": [], "output": None}
      for layer in layer_extractor:
          if isinstance(layer[0], MessagePassing) or isinstance(
              layer[0], GNNPool
          ):
              if isinstance(layer[0], GNNPool):
                  pool_flag = True
              if step["module"] and step["input"] is not None:
                  walk_steps.append(step)
              step = {"input": layer[1], "module": [], "output": None}
          if pool_flag and split_fc and isinstance(layer[0], nn.Linear):
              if step["module"]:
                  fc_steps.append(step)
              step = {"input": layer[1], "module": [], "output": None}
          step["module"].append(layer[0])
          step["output"] = layer[2]
  
      for walk_step in walk_steps:
          if (
              hasattr(walk_step["module"][0], "nn")
              and walk_step["module"][0].nn is not None
          ):
              # We don't allow any outside nn during message flow process in GINs
              walk_step["module"] = [walk_step["module"][0]]
  
      if split_fc:
          if step["module"]:
              fc_steps.append(step)
          return walk_steps, fc_steps
      else:
          fc_step = step
  
      return walk_steps, fc_step
  
imports: from typing import List, Tuple, Dict
  from math import sqrt
  import torch
  from torch import Tensor
  import torch.nn
  import matplotlib.pyplot
  import networkx
  from torch_geometric.nn import MessagePassing
  from torch_geometric.utils.loop import add_self_loops, remove_self_loops
  from torch_geometric.data import Data, Batch
  from torch_geometric.utils import to_networkx
  from ..models.utils import subgraph
  from rdkit import Chem
  from matplotlib.axes import Axes
  import numpy
  from ..models import GNNPool
label_desc: Use image node_id 2 for calling the WalkBase obj's underlying member method code with example usage: obj.extract_step(x, edge_index, detach, split_fc) and returns: walk_steps,  fc_step, walk_steps,  fc_steps