self.bert = BertModel.from_pretrained(model_name, output_hidden_states=True)
        self.n_layers = n_layers
        hidden_size = 768 * n_layers
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(hidden_size, 3)