src/transformers/models/bert/modeling_flax_bert.py/FlaxBertPreTrainingHeads/setup
class FlaxBertPreTrainingHeads:    def setup(self):
        self.predictions = FlaxBertLMPredictionHead(self.config, dtype=self.dtype)
        self.seq_relationship = nn.Dense(2, dtype=self.dtype)