class_docstr: This module can be used to attach combination of {Linear, BatchNorm, Relu, Dropout}
  layers and they are fully configurable from the config file. The module also supports
  stacking multiple MLPs.
  
  Examples:
      Linear
      Linear -> BN
      Linear -> ReLU
      Linear -> Dropout
      Linear -> BN -> ReLU -> Dropout
      Linear -> ReLU -> Dropout
      Linear -> ReLU -> Linear -> ReLU -> ...
      Linear -> Linear -> ...
      ...
  
  Accepts a 2D input tensor. Also accepts 4D input tensor of shape `N x C x 1 x 1`.
class_docstr_tok: ["This","module","can","be","used","to","attach","combination","of","{","Linear",",","BatchNorm",",","Relu",",","Dropout","}","layers","and","they","are","fully","configurable","from","the","config","file",".","The","module","also","supports","stacking","multiple","MLPs",".","Examples",":","Linear","Linear","-",">","BN","Linear","-",">","ReLU","Linear","-",">","Dropout","Linear","-",">","BN","-",">","ReLU","-",">","Dropout","Linear","-",">","ReLU","-",">","Dropout","Linear","-",">","ReLU","-",">","Linear","-",">","ReLU","-",">","...","Linear","-",">","Linear","-",">","...","...","Accepts","a","2D","input","tensor",".","Also","accepts","4D","input","tensor","of","shape","`","N","x","C","x","1","x","1","`","."]
docstr: Args:
      batch (torch.Tensor): 2D torch tensor or 4D tensor of shape `N x C x 1 x 1`
  Returns:
      out (torch.Tensor): 2D output torch tensor
docstr_tok: ["Args",":","batch","(","torch.Tensor",")",":","2D","torch","tensor","or","4D","tensor","of","shape","`","N","x","C","x","1","x","1","`","Returns",":","out","(","torch.Tensor",")",":","2D","output","torch","tensor"]
code: def forward(self, batch: torch.Tensor):
      """
      Args:
          batch (torch.Tensor): 2D torch tensor or 4D tensor of shape `N x C x 1 x 1`
      Returns:
          out (torch.Tensor): 2D output torch tensor
      """
      if isinstance(batch, list):
          assert (
              len(batch) == 1
          ), "MLP input should be either a tensor (2D, 4D) or list containing 1 tensor."
          batch = batch[0]
      if batch.ndim > 2:
          assert all(
              d == 1 for d in batch.shape[2:]
          ), f"MLP expected 2D input tensor or 4D tensor of shape NxCx1x1. got: {batch.shape}"
          batch = batch.reshape((batch.size(0), batch.size(1)))
  
      out = self.clf(batch)
      return out
  
file_path: vissl/vissl/models/heads/mlp.py
imports: from typing import List
  import torch
  import torch.nn
  from vissl.config import AttrDict
  from vissl.models.heads import register_model_head
  from vissl.utils.fsdp_utils import fsdp_auto_wrap_bn, fsdp_wrapper
label_desc: Use image node_id 4 for calling the MLP obj's underlying member method code with example usage: obj.forward(batch) and returns: out