src/transformers/models/albert/modeling_flax_albert.py/FlaxAlbertLayers/setup
class FlaxAlbertLayers:    def setup(self):
        self.albert_layers = FlaxAlbertLayer(self.config, name=self.layer_index, dtype=self.dtype)