pytorch - 如何使用条件匹配两个不同 pytorch 张量中的值

Question

我有 4 个张量

ib0 = torch.tensor([14,17])
bb0 = torch.tensor([0,1])
ib1 = torch.tensor([17,4,2])
bb1 = torch.tensor([0,0,1])

我需要根据 bb0 和 bb1 的值获得另一个与 ib0 和 ib1 中的值匹配的张量。我正在寻找的结果是

torch.tensor([[14,17],[14,4],[17,2]])

因为 ib0 中 0 的索引（来自 bb0）的值为 14，而 0 的索引（来自 bb1）在 ib1 中的值为 17,4，对于 bb0 中的 1 的索引也类似。这在波束搜索算法中需要将来自同一波束的索引合并在一起。bb0 和 bb1 张量是梁索引，而 ib0、ib1 是 2 个梁 0 和 1 中的索引。所以上述过程只是确保来自不同梁的索引不合并的一种方法。

标签： pytorch