def forward(self,
                x: torch.Tensor,
                media: torch.Tensor,
                media_locations: Optional[torch.Tensor] = None,
                attend_previous: bool = True):
        """Forward function for perceiver sampler.

        Args:
            x (torch.Tensor): text features of shape (B, T_txt, D_txt).
            media (torch.Tensor): image features of shape
                (B, T_img, n, D_img) where n is the dim of the latents.
            media_locations (torch.Tensor, optional): boolean mask identifying
                the media tokens in x of shape (B, T_txt). Defaults to None.
            attend_previous (bool): If false, ignores immediately preceding
                image and starts attending when following image.
                Defaults to True.
        """
        x = (
            self.attn(
                x,
                media,
                media_locations=media_locations,
                attend_previous=attend_previous,
            ) * self.attn_gate.tanh() + x)
        x = self.ff(x) * self.ff_gate.tanh() + x

        return x
Forward function for perceiver sampler.

        Args:
            x (torch.Tensor): text features of shape (B, T_txt, D_txt).
            media (torch.Tensor): image features of shape
                (B, T_img, n, D_img) where n is the dim of the latents.
            media_locations (torch.Tensor, optional): boolean mask identifying
                the media tokens in x of shape (B, T_txt). Defaults to None.
            attend_previous (bool): If false, ignores immediately preceding
                image and starts attending when following image.
                Defaults to True.
        
mmpretrain/models/multimodal/flamingo/modules.py