From f2b300df6bd46ad16580f0313bc4b30ddde8515d Mon Sep 17 00:00:00 2001 From: LysandreJik Date: Thu, 8 Aug 2019 10:38:57 -0400 Subject: [PATCH] fix #976 --- pytorch_transformers/modeling_bert.py | 4 ++-- pytorch_transformers/modeling_xlm.py | 2 +- pytorch_transformers/modeling_xlnet.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pytorch_transformers/modeling_bert.py b/pytorch_transformers/modeling_bert.py index e13b3d01f9..34eac7f26f 100644 --- a/pytorch_transformers/modeling_bert.py +++ b/pytorch_transformers/modeling_bert.py @@ -925,7 +925,7 @@ class BertForSequenceClassification(BertPreTrainedModel): r""" **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``: Labels for computing the sequence classification/regression loss. - Indices should be in ``[0, ..., config.num_labels]``. + Indices should be in ``[0, ..., config.num_labels - 1]``. If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss), If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy). @@ -1099,7 +1099,7 @@ class BertForTokenClassification(BertPreTrainedModel): r""" **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``: Labels for computing the token classification loss. - Indices should be in ``[0, ..., config.num_labels]``. + Indices should be in ``[0, ..., config.num_labels - 1]``. Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs: **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``: diff --git a/pytorch_transformers/modeling_xlm.py b/pytorch_transformers/modeling_xlm.py index 5acb20e04c..941c8dda2f 100644 --- a/pytorch_transformers/modeling_xlm.py +++ b/pytorch_transformers/modeling_xlm.py @@ -784,7 +784,7 @@ class XLMForSequenceClassification(XLMPreTrainedModel): r""" **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``: Labels for computing the sequence classification/regression loss. - Indices should be in ``[0, ..., config.num_labels]``. + Indices should be in ``[0, ..., config.num_labels - 1]``. If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss), If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy). diff --git a/pytorch_transformers/modeling_xlnet.py b/pytorch_transformers/modeling_xlnet.py index b6adc4de43..e9e75e3ab7 100644 --- a/pytorch_transformers/modeling_xlnet.py +++ b/pytorch_transformers/modeling_xlnet.py @@ -1075,7 +1075,7 @@ class XLNetForSequenceClassification(XLNetPreTrainedModel): r""" **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``: Labels for computing the sequence classification/regression loss. - Indices should be in ``[0, ..., config.num_labels]``. + Indices should be in ``[0, ..., config.num_labels - 1]``. If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss), If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).