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

        x_block, thw_shape_new = self.attn(self.norm1(x), thw_shape)
        atn = AttentionPool(
            pool=self.pool_skip,
            thw_shape=thw_shape,
            has_cls_embed=self.has_cls_embed
        )
        x_res, _ = atn(x)
        x = x_res + self.drop_path(x_block)
        x_norm = 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 (`torch.Tensor`):
Input tensor.
thw_shape (`List`):
The shape of the input tensor (before flattening).
towhee/models/layers/multi_scale_transformer_block.py