examples/research_projects/movement-pruning/emmental/modeling_bert_masked.py/BertLayer/forward
class BertLayer:    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        threshold=None,
    ):
        self_attention_outputs = self.attention(hidden_states, attention_mask, head_mask, threshold=threshold)
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        if self.is_decoder and encoder_hidden_states is not None:
            cross_attention_outputs = self.crossattention(
                attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask
            )
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights

        intermediate_output = self.intermediate(attention_output, threshold=threshold)
        layer_output = self.output(intermediate_output, attention_output, threshold=threshold)
        outputs = (layer_output,) + outputs
        return outputs