class_docstr: A GNN prediction head for edge-level/link-level prediction tasks.
  
  Args:
      dim_in (int): The input feature dimension.
      dim_out (int): The output feature dimension.
class_docstr_tok: ["A","GNN","prediction","head","for","edge-level\/link-level","prediction","tasks",".","Args",":","dim_in","(","int",")",":","The","input","feature","dimension",".","dim_out","(","int",")",":","The","output","feature","dimension","."]
code: def _apply_index(self, batch):
      index = f"{batch.split}_edge_index"
      label = f"{batch.split}_edge_label"
      return batch.x[batch[index]], batch[label]
  
imports: import torch
  import torch_geometric.graphgym.register
  from torch_geometric.graphgym.config import cfg
  from torch_geometric.graphgym.models.layer import MLP, new_layer_config
  from torch_geometric.graphgym.register import register_head
label_desc: Use image node_id 2 for calling the GNNEdgeHead obj's underlying member method code with example usage: obj._apply_index(batch) and returns: batch,  batch