我有一批维度为 [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]]
我在这里错过了什么?