def forward(self, x: torch.Tensor, latents: torch.Tensor):
        """Forward function.

        Args:
            x (torch.Tensor): image features of shape (b, T, n1, D).
            latent (torch.Tensor): latent features of shape (b, T, n2, D).
        """
        x = self.norm_media(x)
        latents = self.norm_latents(latents)

        h = self.heads

        q = self.to_q(latents)
        kv_input = torch.cat((x, latents), dim=-2)
        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
        q = rearrange(q, 'b t n (h d) -> b h t n d', h=h)
        k = rearrange(k, 'b t n (h d) -> b h t n d', h=h)
        v = rearrange(v, 'b t n (h d) -> b h t n d', h=h)
        q = q * self.scale

        # attention
        sim = einsum('... i d, ... j d  -> ... i j', q, k)
        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
        attn = sim.softmax(dim=-1)

        out = einsum('... i j, ... j d -> ... i d', attn, v)
        out = rearrange(out, 'b h t n d -> b t n (h d)', h=h)
        return self.to_out(out)
Forward function.

        Args:
            x (torch.Tensor): image features of shape (b, T, n1, D).
            latent (torch.Tensor): latent features of shape (b, T, n2, D).
        
mmpretrain/models/multimodal/flamingo/modules.py