Python 如何在两个不同大小的PyTorch张量中得到相等元素的索引？

假设我有两个PyTorch张量：

t_1d = torch.Tensor([6, 5, 1, 7, 8, 4, 7, 1, 0, 4, 11, 7, 4, 7, 4, 1])
t = torch.Tensor([4, 7])

我想得到集合之间精确匹配交集的指数 带张量t的张量t_1d

t\u 1d

和

t

的期望输出：

[5,12]

（精确交点的第一个索引）

---

最好在GPU上使用大张量，这样就不会出现循环或Numpy强制转换。

一般来说，我们可以检查

t

中的每个元素在哪里等于

t\u 1d

中的元素

在此之后，将最后一个元素向后移动与第一个元素相同的位置（一般情况下，此处按

-1

）并检查数组是否相等：

intersection = (t_1d == t[0]) & torch.roll(t_1d == t[1], shifts=-1)
torch.where(intersection)[0] # torch.tensor([5, 12])