update new token classification model

511bce58 · thomwolf · 258eb500 · 511bce58
Commit 511bce58 authored Nov 30, 2018 by thomwolf
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

pytorch_pretrained_bert/modeling.py pytorch_pretrained_bert/modeling.py +2 -2

No files found.
--- a/pytorch_pretrained_bert/modeling.py
+++ b/pytorch_pretrained_bert/modeling.py
@@ -932,8 +932,8 @@ class BertForTokenClassification(PreTrainedBertModel):
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
        sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
-        pooled_output = self.dropout(sequence_output)
+        sequence_output = self.dropout(sequence_output)
-        logits = self.classifier(pooled_output)
+        logits = self.classifier(sequence_output)
        if labels is not None:
            loss_fct = CrossEntropyLoss()