Fixed target_mapping preparation for XLNet when batch size > 1 (incl. beam search) (#7267)

39062d05 · guillaume-be · GitHub · 4b3e55bd · 39062d05
Unverified Commit 39062d05 authored Sep 21, 2020 by guillaume-be Committed by GitHub Sep 21, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/modeling_xlnet.py src/transformers/modeling_xlnet.py +1 -1

No files found.
--- a/src/transformers/modeling_xlnet.py
+++ b/src/transformers/modeling_xlnet.py
@@ -1313,7 +1313,7 @@ class XLNetLMHeadModel(XLNetPreTrainedModel):
        target_mapping = torch.zeros(
            (effective_batch_size, 1, sequence_length), dtype=torch.float, device=input_ids.device
        )
-        target_mapping[0, 0, -1] = 1.0
+        target_mapping[:, 0, -1] = 1.0

        inputs = {
            "input_ids": input_ids,