检查张量的每个元素是否包含在列表中

Question

假设我有一个张量A和一个值容器vals。是否有一种干净的方法来返回A与每个元素形状相同的布尔张量，即该元素是否A包含在其中vals？例如：

A = torch.tensor([[1,2,3],
                  [4,5,6]])
vals = [1,5]

# Desired output
torch.tensor([[True,False,False],
              [False,True,False]])

Answer 1

您可以通过 for 循环实现此目的：

sum(A==i for i in B).bool()