    def forward(self, tensor: Tensor, thw_shape: List[int]) -> Tuple[Tensor, List[int]]:
        """
        Args:
            tensor (Tensor): Input tensor.
            thw_shape (List): The shape of the input tensor (before flattening).

        Returns:
            tensor (Tensor): Input tensor after pool.
            thw_shape (List[int]): Output tensor shape (before flattening).
        """
        if not self.has_pool:
            return tensor, thw_shape

        tensor, tensor_dim = _unsqueeze_dims(tensor)

        assert not self.has_cls_embed

        if self.has_cls_embed:
            # TODO: enable has_cls_embed

            # cls_tok: Tensor = torch.tensor(0)  # For typing/torchscriptability
            # if self.has_cls_embed:
            #    cls_tok, tensor = tensor[:, :, :1, :], tensor[:, :, 1:, :]
            raise NotImplementedError("Unsupported the input tensor contains cls token")

        # input shape: B, num_heads, seqlen, head_dim
        B, N, L, C = get_shape(tensor)
        T, H, W = thw_shape
        tensor = ops.reshape()(tensor, [B * N, -1, H, W, C])

        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 = ops.elementwise(FuncEnum.GELU)(tensor)

        tensor = self.pool(tensor)

        shape = get_shape(tensor)
        thw_shape = [shape[1], shape[2], shape[3]]
        L_pooled = shape[1] * shape[2] * shape[3]
        tensor = ops.reshape()(tensor, [B, N, L_pooled, C])

        if self.has_norm and not self.norm_before_pool:
            tensor = self.norm(tensor)

        tensor = _squeeze_dims(tensor, tensor_dim)

        return tensor, thw_shape
Args:
tensor (Tensor): Input tensor.
thw_shape (List): The shape of the input tensor (before flattening).

Returns:
tensor (Tensor): Input tensor after pool.
thw_shape (List[int]): Output tensor shape (before flattening).
python/aitemplate/frontend/nn/multiscale_attention.py