class_docstr: MLP with linear output
class_docstr_tok: ["MLP","with","linear","output"]
code: def forward(self, x):
      if self.linear_or_not:
          # If linear model
          return self.linear(x)
      else:
          # If MLP
          h = x
          for i in range(self.num_layers - 1):
              h = F.relu(self.batch_norms[i](self.linears[i](h)))
          return self.linears[-1](h)
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from dgl.nn.pytorch.conv import GraphConv
  from dgl.nn.pytorch.glob import SortPooling
  from .None import register_model
  from .base import BaseAutoModel
  from ....utils import get_logger
label_desc: Use image node_id 2 for calling the MLP obj's underlying member method code with example usage: obj.forward(x) and returns: self, self