examples/research_projects/bert-loses-patience/pabee/modeling_pabee_bert.py/BertEncoderWithPabee/adaptive_forward
class BertEncoderWithPabee:    def adaptive_forward(self, hidden_states, current_layer, attention_mask=None, head_mask=None):
        layer_outputs = self.layer[current_layer](hidden_states, attention_mask, head_mask[current_layer])

        hidden_states = layer_outputs[0]

        return hidden_states