I would like to create a mask for a tensor based on the existence of another tensor within the first. This is easy to do if I am checking for a scalar, but I could not find a way to do it tensor-wise in the way I need.

What I can do element-wise:

x = torch.Tensor([[1,2,3,4],[4,3,2,1]])
mask = ( x == 2 )

mask equals tensor([[False, True, False, False], [False, False, True, False]]).

This is what I would like to do:

x = torch.Tensor([[1,2,3,4],[4,3,2,1]])
mask = ( x == torch.Tensor([1,2]) ) #this does not actually work

Which would have mask equal tensor([[True, True, False, False], [ False, False, False, False]]).

Any help is appreciated, thanks!