def forward(
        self, x: Tensor, t_shape: int, h_shape: int, w_shape: int
    ) -> Tuple[Tensor, List[int]]:
        """
        Args:
            x (Tensor): Input tensor.
            thw_shape (List): The shape of the input tensor (before flattening).
        """
        thw_shape = [t_shape, h_shape, w_shape]
        x_norm = (
            ait_ncl2nlc(self.norm1(ait_ncl2nlc(x)))
            if self.norm1_is_batchnorm_1d
            else self.norm1(x)
        )
        x_block, thw_shape_new = self.attn(x_norm, thw_shape)
        x_res, _ = self._attention_pool(x, thw_shape)
        x = x_res + self.drop_path(x_block)
        x_norm = (
            ait_ncl2nlc(self.norm2(ait_ncl2nlc(x)))
            if self.norm2_is_batchnorm_1d
            else self.norm2(x)
        )
        x_mlp = self.mlp(x_norm)
        if self.dim != self.dim_out:
            x = self.proj(x_norm)
        x = x + self.drop_path(x_mlp)

        return x, thw_shape_new

        Args:
            x (Tensor): Input tensor.
            thw_shape (List): The shape of the input tensor (before flattening).
        
python/aitemplate/frontend/nn/multiscale_attention.py