对pytorch 矩阵中的元素排序并返回对应索引

可以使用`torch.sort()`函数进行排序，它会返回排序后的值和对应的索引。示例代码如下：

import torch

# 创建一个 2x3 的矩阵
matrix = torch.tensor([[4, 2, 3], [1, 5, 6]])

# 对每行元素进行排序并返回对应的索引
sorted_values, sorted_indices = torch.sort(matrix, dim=1)

print("排序后的结果：")
print(sorted_values)
print("对应的索引：")
print(sorted_indices)

输出结果为：

排序后的结果：
tensor([[2, 3, 4],
        [1, 5, 6]])
对应的索引：
tensor([[1, 2, 0],
        [0, 1, 2]])

在上面的代码中，我们使用了`dim=1`来表示对每行元素进行排序。如果要对每列元素进行排序，只需要将`dim`参数设置为`0`即可。