torch argsort根据dim=...返回行或列内的索引
我怎么才能得到一个二维索引...

[[ r1,r2,r3,...], [c1,c2,c3,.....]]

谢谢...这是我做的

#https://github.com/pytorch/pytorch/issues/35674
def unravel_indices(indices, shape):
    coord = []

    for dim in reversed(shape):
        coord.append(torch.fmod(indices, dim))
        indices = torch.div(indices, dim, rounding_mode='floor')

    coord = torch.stack(coord[::-1], dim=-1)

    return coord
torch.unravel_indices = unravel_indices