Skip to content

Commit fe9051c

Browse files
urialonvince62s
authored andcommitted
Fix a potential IndexError when translating with replace_unk (#1469)
* Fix IndexError which happens with replace_unk, when the argmax of the attention is on the padding instead of a real source token
1 parent bed1788 commit fe9051c

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

onmt/translate/translation.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def _build_target_tokens(self, src, src_vocab, src_raw, pred, attn):
4747
if self.replace_unk and attn is not None and src is not None:
4848
for i in range(len(tokens)):
4949
if tokens[i] == tgt_field.unk_token:
50-
_, max_index = attn[i].max(0)
50+
_, max_index = attn[i][:len(src_raw)].max(0)
5151
tokens[i] = src_raw[max_index.item()]
5252
if self.phrase_table != "":
5353
with open(self.phrase_table, "r") as f:

0 commit comments

Comments
 (0)