S = self.height * self.width
Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S]
f_Bs = answers
- return (B_s == f_Bs).long().min(dim=-1).values > 0
+ return (Bs == f_Bs).long().min(dim=-1).values > 0
def generate_prompts_and_answers(self, nb, tasks=None, device="cpu"):
if tasks is None: