我有两个张量,张量a和张量b
我想得到张量b中所有值的索引
比如说
a = torch.Tensor([1,2,2,3,4,4,4,5])
b = torch.Tensor([1,2,4])
我想要张量a中的1, 2, 4
的索引。我可以通过以下代码来实现这一点
a = torch.Tensor([1,2,2,3,4,4,4,5])
b = torch.Tensor([1,2,4])
mask = torch.zeros(a.shape).type(torch.bool)
print(mask)
for e in b:
mask = mask + (a == e)
print(mask)
没有for
我怎么做
如果您不想使用for循环,可以使用列表理解:
如果你甚至不想使用“for”这个词,你可以将张量转换成numpy,并使用numpy索引
更新
可能误解了你的问题。在这种情况下,我想说实现这一点的最佳方法是通过列表理解。(切片可能无法实现这一点
它迭代a中的每个元素,获取它们的索引和值,如果值在b中,则获取索引
这就是你想要的吗
将导致:
相关问题 更多 >
编程相关推荐