Bug fix

app.py

@@ -56,11 +56,10 @@ class Processor:
 
         prompts = [question + (f' \\n Knowledge: {knowledge} \\n Answer: ' if knowledge != '' else ' \\n Answer:') for knowledge in knowledges]
         prompts_tok = self.tokenizer(prompts, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_question_len + max_knowledge_len).to(device) # (1+K, QL+KL)
-        print(prompts_tok.input_ids.size(), choices_ids[0].unsqueeze(0).expand(len(knowledges), -1).size())
         output = self.model(
             input_ids=prompts_tok.input_ids,
             attention_mask=prompts_tok.attention_mask,
-            labels=choices_ids[0].unsqueeze(0).expand(len(knowledges), -1),
+            labels=choices_ids[0].unsqueeze(0).repeat(len(knowledges), 1),
         )
         logitsss = output.logits # (1+K, AL, V)
         logitss = logitsss[:, 0, :] # (1+K, V)