src/transformers/models/beit/modeling_flax_beit.py/FlaxBeitPooler/setup
class FlaxBeitPooler:    def setup(self):
        if self.config.use_mean_pooling:
            self.layernorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)