src/transformers/models/albert/modeling_flax_albert.py/FlaxAlbertForPreTrainingModule/setup
class FlaxAlbertForPreTrainingModule:    def setup(self):
        self.albert = FlaxAlbertModule(config=self.config, dtype=self.dtype)
        self.predictions = FlaxAlbertOnlyMLMHead(config=self.config, dtype=self.dtype)
        self.sop_classifier = FlaxAlbertSOPHead(config=self.config, dtype=self.dtype)