pytorch 中的索引和切片

问题描述 投票:0回答:0

我有一批维度为 [B,n,n] 的二维张量和维度为 [B,2] 的坐标张量。

a = torch.arange(48).reshape((3,4,4))
coords = torch.tensor([[0,1],[1,2],[1,3]],dtype=int)

给我想要的结果的代码是:

a[torch.arange(3),coords[:,0],coords[:,1]]

我只是不明白为什么我不能只使用以下内容,因为我认为“:”意味着采用所有索引:

a[:,coords[:,0],coords[:,1]]

我在这里错过了什么?

python indexing pytorch slice
© www.soinside.com 2019 - 2024. All rights reserved.