如何根据Pytorch中列表的值定义掩码函数

我想根据张量的值来屏蔽它。在下面的函数中,如果我传递一个范围(第二部分),它可以工作,但我希望有一个包含各种值的列表 prompt_ids

RuntimeError: Boolean value of Tensor with more than one value is ambiguous

职能:

   def get_prompt_token_fn(self):
        if self.prompt_ids:
            return lambda x: x in self.prompt_ids
        else:
            return lambda x: (x>=self.id_offset)&(x<self.id_offset+self.length)

问题是什么?我如何解决?