src/transformers/models/bert/modeling_flax_bert.py/FlaxBertOnlyMLMHead/setup
class FlaxBertOnlyMLMHead:    def setup(self):
        self.predictions = FlaxBertLMPredictionHead(self.config, dtype=self.dtype)