torch.take_along_dim¶
- torch.take_along_dim(input, indices, dim=None, *, out=None) Tensor ¶
从
input
中沿给定的dim
选择来自indices
中的一维索引处的数值。如果
dim
为 None,则输入数组将被视为已扁平化为一维。返回沿维度索引的函数,例如
torch.argmax()
和torch.argsort()
,旨在与该函数一起使用。请参见下面的示例。注意
该函数类似于 NumPy 的 take_along_axis。另请参见
torch.gather()
。- 参数
- 关键字参数
out (Tensor, 可选) – 输出张量。
示例
>>> t = torch.tensor([[10, 30, 20], [60, 40, 50]]) >>> max_idx = torch.argmax(t) >>> torch.take_along_dim(t, max_idx) tensor([60]) >>> sorted_idx = torch.argsort(t, dim=1) >>> torch.take_along_dim(t, sorted_idx, dim=1) tensor([[10, 20, 30], [40, 50, 60]])