docstr: Args:
      tensor (torch.Tensor): Input tensor.
      thw_shape (List): The shape of the input tensor (before flattening).
  
  Returns:
      tensor (torch.Tensor): Input tensor after pool.
      thw_shape (List[int]): Output tensor shape (before flattening).
docstr_tok: ["Args",":","tensor","(","torch.Tensor",")",":","Input","tensor",".","thw_shape","(","List",")",":","The","shape","of","the","input","tensor","(","before","flattening",")",".","Returns",":","tensor","(","torch.Tensor",")",":","Input","tensor","after","pool",".","thw_shape","(","List","[","int","]",")",":","Output","tensor","shape","(","before","flattening",")","."]
code: def forward(
      self, tensor: torch.Tensor, thw_shape: List[int]
  ) -> Tuple[torch.Tensor, List[int]]:
      """
      Args:
          tensor (torch.Tensor): Input tensor.
          thw_shape (List): The shape of the input tensor (before flattening).
  
      Returns:
          tensor (torch.Tensor): Input tensor after pool.
          thw_shape (List[int]): Output tensor shape (before flattening).
      """
      if not self.has_pool:
          return tensor, thw_shape
      tensor_dim = tensor.ndim
  
      if torch.jit.is_scripting():
          tensor, tensor_dim = _unsqueeze_dims_jit(tensor)
      else:
          tensor, tensor_dim = _unsqueeze_dims_fx(tensor)
  
      cls_tok: torch.Tensor = torch.tensor(
          0
      )  # For typing/torchscriptability
      if self.has_cls_embed:
          cls_tok, tensor = tensor[:, :, :1, :], tensor[:, :, 1:, :]
  
      B, N, L, C = tensor.shape
      T, H, W = thw_shape
      tensor = (
          tensor.reshape(B * N, T, H, W, C)
          .permute(0, 4, 1, 2, 3)
          .contiguous()
      )
  
      if self.norm_before_pool:
          # If use BN, we apply norm before pooling instead of after pooling.
          tensor = self.norm(tensor)
          # We also empirically find that adding a GELU here is beneficial.
          tensor = torch.nn.functional.gelu(tensor)
  
      tensor = self.pool(tensor)
  
      thw_shape = [tensor.shape[2], tensor.shape[3], tensor.shape[4]]
      L_pooled = tensor.shape[2] * tensor.shape[3] * tensor.shape[4]
      tensor = tensor.reshape(B, N, C, L_pooled).transpose(2, 3)
      if self.has_cls_embed:
          tensor = torch.cat((cls_tok, tensor), dim=2)
      if self.has_norm and not self.norm_before_pool:
          tensor = self.norm(tensor)
  
      if torch.jit.is_scripting():
          tensor = _squeeze_dims_jit(tensor, tensor_dim)
      else:
          tensor = _squeeze_dims_fx(tensor, tensor_dim)
  
      return tensor, thw_shape
  
imports: from typing import Callable, List, Optional, Tuple
  import numpy
  import torch
  import torch.nn
  from torch.nn.common_types import _size_3_t
  from .drop_path import DropPath
label_desc: Use image node_id 2 for calling the _AttentionPool obj's underlying member method code with example usage: obj.forward(tensor, thw_shape) and returns: tensor,  thw_shape, tensor,  thw_shape