How to find element in tensor A that is in tensor b

Given tensor A=torch.tensor([[2, 1, 5], [3, 2, 4]]) and B=torch.tensor([1, 3, 5])
How can I return tensor C that tells whether element in A is in B, i.e. C=torch.tensor([[False, True, True], [True, False, False]]).

2 Likes

You can do this with a for loop:

>>> sum(A==i for i in B).bool()
tensor([[False,  True,  True],
        [ True, False, False]])
1 Like