class_docstr: TODO
docstr: Args:
      batch (torch.Tensor): 3D torch tensor.
  Returns:
      out (torch.Tensor): 2D output torch tensor
docstr_tok: ["Args",":","batch","(","torch.Tensor",")",":","3D","torch","tensor",".","Returns",":","out","(","torch.Tensor",")",":","2D","output","torch","tensor"]
code: def forward(self, batch: torch.Tensor):
      """
      Args:
          batch (torch.Tensor): 3D torch tensor.
      Returns:
          out (torch.Tensor): 2D output torch tensor
      """
      assert (
          len(batch.shape) == 3
      ), "Expecting shape (batch, seq, embed)"
      batch = batch.permute((0, 2, 1)).unsqueeze(3)
  
      assert (
          len(batch.shape) == 4
      ), "Expecting shape (batch, embed, seq, 1)"
      out = self.channel_bn(batch)
      out = torch.flatten(out, start_dim=2)
  
      # Reshape: batch, embed, seq -> seq, batch, embed
      out = out.permute((2, 0, 1))
      out = self.clf(out)
      return EnsembleOutput(out)
  
file_path: vissl/vissl/models/heads/mlp_pooled.py
imports: from typing import List
  import torch
  import torch.nn
  from vissl.config import AttrDict
  from vissl.losses.cross_entropy_multiple_output_single_target import EnsembleOutput
  from vissl.models.heads import register_model_head
  from vissl.models.heads.mlp import MLP
label_desc: Use image node_id 2 for calling the LinearEvalPooledMLP obj's underlying member method code with example usage: obj.forward(batch) and returns: EnsembleOutput