class_docstr: General GNN model: encoder + stage + head
  
  Args:
      dim_in (int): Input dimension
      dim_out (int): Output dimension
      **kwargs (optional): Optional additional args
code: def __init__(self, dim_in, dim_out, **kwargs):
      super(GNN, self).__init__()
      GNNStage = stage_dict[cfg.gnn.stage_type]
      GNNHead = head_dict[cfg.dataset.task]
  
      self.encoder = FeatureEncoder(dim_in)
      dim_in = self.encoder.dim_in
  
      if cfg.gnn.layers_pre_mp > 0:
          self.pre_mp = GNNPreMP(dim_in, cfg.gnn.dim_inner)
          dim_in = cfg.gnn.dim_inner
      if cfg.gnn.layers_mp > 0:
          self.mp = GNNStage(
              dim_in=dim_in,
              dim_out=cfg.gnn.dim_inner,
              num_layers=cfg.gnn.layers_mp,
          )
      self.post_mp = GNNHead(dim_in=cfg.gnn.dim_inner, dim_out=dim_out)
  
      self.apply(init_weights)
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  import graphgym.register
  from graphgym.config import cfg
  from graphgym.init import init_weights
  from graphgym.models.feature_encoder_pyg import edge_encoder_dict, node_encoder_dict
  from graphgym.models.head_pyg import head_dict
  from graphgym.models.layer_pyg import BatchNorm1dEdge, BatchNorm1dNode, GeneralLayer, GeneralMultiLayer
label_desc: Use image node_id 1 to create a new GNN object from inherited base classes: nn with example: obj = GNN(dim_in, dim_out)