Fix merge issue

yifan1130 · Sep 22, 2023 · ef0d437 · ef0d437
1 parent 0ddab2b
commit ef0d437
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/eval/utils.py b/eval/utils.py
@@ -113,8 +113,6 @@ def get_next_word_predictions(model, tokenizer, prompts, candidate_token_ids=Non
             attention_mask = attention_mask.cuda()
 
         batch_logits = model(input_ids=batch_input_ids, attention_mask=attention_mask).logits[:, -1, :]
-        if candidate_token_ids is not None:
-            batch_logits = batch_logits[:, candidate_token_ids]
         batch_probs = torch.softmax(batch_logits, dim=-1)
         if candidate_token_ids is not None:
             batch_probs = batch_probs[:, candidate_token_ids]