src/transformers/models/big_bird/modeling_flax_big_bird.py/FlaxBigBirdPreTrainingHeads/__call__
class FlaxBigBirdPreTrainingHeads:    def __call__(self, hidden_states, pooled_output, shared_embedding=None):
        prediction_scores = self.predictions(hidden_states, shared_embedding=shared_embedding)
        seq_relationship_score = self.seq_relationship(pooled_output)
        return prediction_scores, seq_relationship_score