src/transformers/models/roberta/modeling_flax_roberta.py/FlaxRobertaForMultipleChoiceModule/setup
class FlaxRobertaForMultipleChoiceModule:    def setup(self):
        self.roberta = FlaxRobertaModule(
            config=self.config,
            dtype=self.dtype,
            gradient_checkpointing=self.gradient_checkpointing,
        )
        self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob)
        self.classifier = nn.Dense(1, dtype=self.dtype)