From 1ebfeb79469d544a2bd817aa32c77e0514485ff9 Mon Sep 17 00:00:00 2001 From: Lysandre Date: Tue, 4 Feb 2020 15:56:16 -0500 Subject: [PATCH] Cast to long when masking tokens --- examples/run_lm_finetuning.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/run_lm_finetuning.py b/examples/run_lm_finetuning.py index c37af5ef48..e7a09abd94 100644 --- a/examples/run_lm_finetuning.py +++ b/examples/run_lm_finetuning.py @@ -195,6 +195,7 @@ def _rotate_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) - def mask_tokens(inputs: torch.Tensor, tokenizer: PreTrainedTokenizer, args) -> Tuple[torch.Tensor, torch.Tensor]: """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """ + inputs = inputs.clone().type(dtype=torch.long) labels = inputs.clone() # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa) probability_matrix = torch.full(labels.shape, args.mlm_probability)