// were words. Add None in place of any ignored <unk>"s.
if self._ignore_unk:
seq_word_ids = word_ids[:, seq_index]
logprob_iter = iter(seq_logprobs)
seq_logprobs = [None if word_id == self._unk_id
else next(logprob_iter)
for word_id in seq_word_ids[1:]]
result.append(seq_logprobs)
After Change
seq_logprobs = logprobs[seq_mask == 1, seq_index]
// The new mask also masks excluded tokens, replace those with None.
seq_mask = new_mask[seq_mask == 1, seq_index]
seq_logprobs = [lp if m == 1 else None
for lp, m in zip(seq_logprobs, seq_mask)]
result.append(seq_logprobs)
return result
