class_docstr: MobileNet-V3 head as done in Torchvision implementation
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
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 or list containing 1 tensor."
          batch = batch[0]
  
      if self.with_bn:
          if len(batch.shape) == 2:
              batch = batch.unsqueeze(2).unsqueeze(3)
          assert len(batch.shape) == 4
          out = self.channel_bn(batch)
      else:
          if len(batch.shape) == 2:
              batch = batch.unsqueeze(2).unsqueeze(3)
          out = batch
  
      out = self.classifier(out)
      return out
  
file_path: vissl/vissl/models/heads/mobilenet_v3_head.py
imports: import torch
  import torch.nn
  from vissl.config import AttrDict
  from vissl.models.heads import register_model_head
label_desc: Use image node_id 2 for calling the MobileNetV3HeadTIMM obj's underlying member method code with example usage: obj.forward(batch) and returns: out