src/transformers/models/albert/modeling_flax_albert.py/FlaxAlbertLayerCollections/setup
class FlaxAlbertLayerCollections:    def setup(self):
        self.albert_layers = FlaxAlbertLayerCollection(self.config, dtype=self.dtype)