小编use*_*788的帖子

我试图索引多维张量中沿最后一个维度的最大元素。例如说我有一个张量

A = torch.randn((5, 2, 3))
_, idx = torch.max(A, dim=2)

idx在这里存储最大索引，看起来可能像

>>>> A
tensor([[[ 1.0503,  0.4448,  1.8663],
     [ 0.8627,  0.0685,  1.4241]],

    [[ 1.2924,  0.2456,  0.1764],
     [ 1.3777,  0.9401,  1.4637]],

    [[ 0.5235,  0.4550,  0.2476],
     [ 0.7823,  0.3004,  0.7792]],

    [[ 1.9384,  0.3291,  0.7914],
     [ 0.5211,  0.1320,  0.6330]],

    [[ 0.3292,  0.9086,  0.0078],
     [ 1.3612,  0.0610,  0.4023]]])
>>>> idx
tensor([[ 2,  2],
    [ 0,  2],
    [ 0,  0],
    [ 0,  2],
    [ 1,  0]])

我希望能够访问这些索引并基于它们分配给另一个张量。意思是我想做

B = torch.new_zeros(A.size())
B[idx] = A[idx]

其中B在所有地方都是0，除了A在最后一个维度上最大。那是B应该存储

>>>>B …