def forward(
        self,
        lang_x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **decoder_layer_kwargs,
    ):
        """Forward function.

        Args:
            lang_x (torch.Tensor): language inputs.
            attention_mask (torch.Tensor, optional): text attention mask.
                Defaults to None.
            **decoder_layer_kwargs: Other decoder layer keyword arguments.
        """
        if self.gated_cross_attn_layer is None:
            return self.decoder_layer(
                lang_x, attention_mask=attention_mask, **decoder_layer_kwargs)

        if self.vis_x is None:
            raise ValueError('vis_x must be conditioned before forward pass')

        if self.media_locations is None:
            raise ValueError(
                'media_locations must be conditioned before forward pass')

        lang_x = self.gated_cross_attn_layer(
            lang_x,
            self.vis_x,
            media_locations=self.media_locations,
            attend_previous=self.attend_previous,
        )
        lang_x = self.decoder_layer(
            lang_x, attention_mask=attention_mask, **decoder_layer_kwargs)
        return lang_x
Forward function.

        Args:
            lang_x (torch.Tensor): language inputs.
            attention_mask (torch.Tensor, optional): text attention mask.
                Defaults to None.
            **decoder_layer_kwargs: Other decoder layer keyword arguments.
        
mmpretrain/models/multimodal/flamingo/modules.py