src/transformers/models/albert/modeling_flax_albert.py/FlaxAlbertSOPHead/setup
class FlaxAlbertSOPHead:    def setup(self):
        self.dropout = nn.Dropout(self.config.classifier_dropout_prob)
        self.classifier = nn.Dense(2, dtype=self.dtype)