How can I define a mask function based on the values of a list in Pytorch

I want to mask a tensor based on its values. In the following function if I pass a range (second part) it works, but I want to have a list with various values prompt_ids (3, 8, 9, 30). But it doesn’t work and throws error.

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)

What’s the problem and how can I resolve it?

I found isin function exists in pytorch 1.10 but not lower versions.
My current solution is something like this, if there are better solutions you can provide:

def isin(ar1):
    return (ar1[..., None] == self.promp_ids).any(-1)

Currently Have on device error (one on cuda and other on cpu)