class_docstr: Head of GNN for node prediction.
code: def _apply_index(self, batch):
      if batch.node_label_index.shape[0] == batch.node_label.shape[0]:
          return batch.x[batch.node_label_index], batch.node_label
      else:
          return (
              batch.x[batch.node_label_index],
              batch.node_label[batch.node_label_index],
          )
  
imports: import torch.nn
  from torch_geometric.graphgym.register import register_head
label_desc: Use image node_id 2 for calling the ExampleNodeHead obj's underlying member method code with example usage: obj._apply_index(batch) and returns: batch,  batch, batch,  batch