src/transformers/models/roberta/modeling_flax_roberta.py/FlaxRobertaModule/setup
class FlaxRobertaModule:    def setup(self):
        self.embeddings = FlaxRobertaEmbeddings(self.config, dtype=self.dtype)
        self.encoder = FlaxRobertaEncoder(self.config, dtype=self.dtype)
        self.pooler = FlaxRobertaPooler(self.config, dtype=self.dtype)