fix bart loss masking (#9131)

This commit is contained in:
Patrick von Platen
2020-12-15 18:17:17 +01:00
committed by GitHub
parent 3caba8d35f
commit 80bdb9c31a
@@ -70,6 +70,11 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int):
index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
prev_output_tokens[:, 1:] = input_ids[:, :-1]
assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
# replace possible -100 values in labels by `pad_token_id`
prev_output_tokens.masked_fill_(prev_output_tokens == -100, pad_token_id)
return prev_output_tokens