Patched error

logits_ngrams.py

@@ -24,7 +24,10 @@ def _no_repeat_ngram_logits(input_ids, cur_len, logits, batch_size=1, no_repeat_
         # from fairseq: https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345
         banned_tokens = _calc_banned_tokens(input_ids, batch_size, no_repeat_ngram_size, cur_len)
         for batch_idx in range(batch_size):
-            logits[batch_idx, [token for token in banned_tokens[batch_idx] if skip_tokens is not None and int(token) not in skip_tokens]] = -float("inf")
+            if skip_tokens is not None:
+                logits[batch_idx, [token for token in banned_tokens[batch_idx] if int(token) not in skip_tokens]] = -float("inf")
+            else:
+                logits[batch_idx, banned_tokens[batch_idx]] = -float("inf")
 
     return logits
 
@@ -35,7 +38,7 @@ def _calc_banned_tokens(prev_input_ids, num_hypos, no_repeat_ngram_size, cur_len
         return [[] for _ in range(num_hypos)]
     generated_ngrams = [{} for _ in range(num_hypos)]
     for idx in range(num_hypos):
-        gen_tokens = prev_input_ids[idx]  # .tolist()
+        gen_tokens = prev_input_ids[idx].tolist()
         generated_ngram = generated_ngrams[idx]
         for ngram in zip(*[gen_tokens[i:] for i in range(no_repeat_ngram_size)]):