setup training

This commit is contained in:
Rémi Louf 2019-11-08 15:48:31 +01:00 committed by Julien Chaumond
parent 4735c2af07
commit 9f75565ea8
2 changed files with 1 additions and 3 deletions

View File

@ -10,5 +10,3 @@ regex
sentencepiece
# For XLM
sacremoses
# For ROUGE
pyrouge

View File

@ -166,7 +166,7 @@ class BeamSearch(object):
for step in range(self.max_length):
decoder_input = fit_to_block_size(self.growing_beams, block_size, self.pad_token_id)
kwargs_decoder["attention_mask"] = build_mask(decoder_input)
kwargs_decoder["attention_mask"] = build_mask(decoder_input, self.pad_token_id)
outputs = self.model.decoder(decoder_input, **kwargs_decoder)
next_token_scores = outputs[0][:, -1, :].squeeze(1)