examples/research_projects/bert-loses-patience/pabee/modeling_pabee_albert.py/AlbertTransformerWithPabee/adaptive_forward
class AlbertTransformerWithPabee:    def adaptive_forward(self, hidden_states, current_layer, attention_mask=None, head_mask=None):
        if current_layer == 0:
            hidden_states = self.embedding_hidden_mapping_in(hidden_states)
        else:
            hidden_states = hidden_states[0]

        layers_per_group = int(self.config.num_hidden_layers / self.config.num_hidden_groups)

        # Index of the hidden group
        group_idx = int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups))

        layer_group_output = self.albert_layer_groups[group_idx](
            hidden_states,
            attention_mask,
            head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group],
        )
        hidden_states = layer_group_output[0]

        return (hidden_states,)