Do not use .T and prefer .t() to be able to create a record in the exported graph.

Signed-off-by: Morgan Funtowicz <morgan@huggingface.co>
This commit is contained in:
Morgan Funtowicz
2020-11-16 16:17:26 +01:00
parent 18b9a8ebf8
commit f85f993442
+1 -1
View File
@@ -1204,7 +1204,7 @@ class BartForSequenceClassification(PretrainedBartModel):
# Attempt to gather the latest eos_token representation in an ONNX compatible way
# (-1: remove the batch indexes from the vector, -1: Take the last occurrence of eos_token)
sentence_representation = torch.index_select(x, 1, eos_mask.nonzero().T[-1, -1])
sentence_representation = torch.index_select(x, 1, eos_mask.nonzero().t()[-1, -1])
logits = self.classification_head(sentence_representation).squeeze(1)
loss = None