Remove pad tokens from attention during decoding

pytorch_translate/beam_decode.py

@@ -171,6 +171,8 @@ def _generate(self, encoder_input, beam_size=None, maxlen=None, prefix_tokens=No
 
         attn = scores.new(bsz * beam_size, src_encoding_len, maxlen + 2)
         attn_buf = attn.clone()
+
+        nonpad_idxs = src_tokens.ne(self.pad)
 
         # list of completed sentences
         finalized = [[] for i in range(bsz)]
@@ -258,11 +260,13 @@ def finalize_hypos(step, bbsz_idx, eos_scores, unfinalized_scores=None):
                 sents_seen.add(sent)
 
                 def get_hypo():
-                    _, alignment = attn_clone[i].max(dim=0)
+                    # remove padding tokens from attn scores
+                    hypo_attn = attn_clone[i][nonpad_idxs[sent]]
+                    _, alignment = hypo_attn.max(dim=0)
                     return {
                         "tokens": tokens_clone[i],
                         "score": score,
-                        "attention": attn_clone[i],  # src_len x tgt_len
+                        "attention": hypo_attn,  # src_len x tgt_len
                         "alignment": alignment,
                         "positional_scores": pos_scores[i],
                     }