class_docstr: A standard Linear classification module that can be attached to several
  layers of the model to evaluate the representation quality of features.
  
   The layers attached are:
      BatchNorm2d -> Linear (1 or more)
  
  Accepts a 4D input tensor. If you want to use 2D input tensor instead,
  use the "mlp" head directly.
class_docstr_tok: ["A","standard","Linear","classification","module","that","can","be","attached","to","several","layers","of","the","model","to","evaluate","the","representation","quality","of","features",".","The","layers","attached","are",":","BatchNorm2d","-",">","Linear","(","1","or","more",")","Accepts","a","4D","input","tensor",".","If","you","want","to","use","2D","input","tensor","instead",",","use","the","``","mlp","''","head","directly","."]
docstr: Args:
      batch (torch.Tensor): 4D torch tensor. This layer is meant to be attached at several
                            parts of the model to evaluate feature representation quality.
                            For 2D input tensor, the tensor is unsqueezed to NxDx1x1 and
                            then eval_mlp is applied
  Returns:
      out (torch.Tensor): 2D output torch tensor
docstr_tok: ["Args",":","batch","(","torch.Tensor",")",":","4D","torch","tensor",".","This","layer","is","meant","to","be","attached","at","several","parts","of","the","model","to","evaluate","feature","representation","quality",".","For","2D","input","tensor",",","the","tensor","is","unsqueezed","to","NxDx1x1","and","then","eval_mlp","is","applied","Returns",":","out","(","torch.Tensor",")",":","2D","output","torch","tensor"]
code: def forward(self, batch: torch.Tensor):
      """
      Args:
          batch (torch.Tensor): 4D torch tensor. This layer is meant to be attached at several
                                parts of the model to evaluate feature representation quality.
                                For 2D input tensor, the tensor is unsqueezed to NxDx1x1 and
                                then eval_mlp is applied
      Returns:
          out (torch.Tensor): 2D output torch tensor
      """
      # in case of a 2D tensor input, unsqueeze the tensor to (N x D x 1 x 1) and apply
      # eval mlp normally.
      if len(batch.shape) == 2:
          batch = batch.unsqueeze(2).unsqueeze(3)
      assert (
          len(batch.shape) == 4
      ), f"Eval MLP head expects 4D tensor input: {batch.shape}"
      out = self.channel_bn(batch)
      out = torch.flatten(out, start_dim=1)
      out = self.clf(out)
      return out
  
file_path: vissl/vissl/models/heads/linear_eval_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.models.heads.mlp import MLP
  from vissl.models.heads.mlp_res import MLPResidual
  from vissl.utils.fsdp_utils import fsdp_auto_wrap_bn, fsdp_wrapper
label_desc: Use image node_id 2 for calling the LinearEvalMLP obj's underlying member method code with example usage: obj.forward(batch) and returns: out