src/transformers/models/distilbert/modeling_flax_distilbert.py/FlaxTransformerEncoder/setup
class FlaxTransformerEncoder:    def setup(self):
        self.layer = FlaxTransformer(self.config, dtype=self.dtype)