def METHOD_NAME(self, encoder_states, encoder_mask):
    """
    Args:
        encoder_states: output of the encoder (B x L_enc x H)
        encoder_mask: encoder inputs mask (B x L_enc)
    """
    # all hidden values are active
    hidden_mask = torch.ones(
        encoder_states.shape[0], self._hidden_steps, dtype=encoder_mask.dtype, device=encoder_mask.device
    )
    # initialize hidden state
    if self._hidden_init_method == "params":
        # initialize latent with learned parameters
        hidden_states = self.init_hidden.unsqueeze(0).expand(encoder_states.shape[0], -1, -1)
        hidden_states = self.init_cross_att(
            decoder_states=hidden_states,
            decoder_mask=hidden_mask,
            encoder_states=encoder_states,
            encoder_mask=encoder_mask,
        )
    elif self._hidden_init_method == "bridge":
        # initialize latent with attention bridge
        hidden_states = self.att_bridge(hidden=encoder_states, hidden_mask=encoder_mask,)
    # apply block (cross-attention, self-attention) multiple times
    # for block in range(self._hidden_blocks):
    for self_att, cross_att in zip(self.self_att_layers, self.cross_att_layers):
        residual = hidden_states
        # cross attention of hidden over encoder states
        hidden_states = cross_att(
            decoder_states=hidden_states,
            decoder_mask=hidden_mask,
            encoder_states=encoder_states,
            encoder_mask=encoder_mask,
        )
        # self-attention over hidden
        hidden_states = self_att(encoder_states=hidden_states, encoder_mask=hidden_mask,)
        # residual connection
        hidden_states += residual
    return hidden_states, hidden_mask