如何按第一维对张量进行排序

我有一个 2D 张量，我想按第一个维度排序，如下例所示：

a = torch.FloatTensor(
   [[5, 5],
    [5, 3],
    [3, 5],
    [6, 4],
    [3, 7]])

我预计排序后的结果是这样的：

a = torch.FloatTensor(
   [[3, 5],
    [3, 7],
    [5, 3],
    [5, 5],
    [6, 4]])

在pytorch中可以做到这一点吗？我知道可以在 numpy 中做到这一点，但我想使用 torch.GPU 在 GPU 中做到这一点。

按第一列排序并使用索引对整个数组进行排序：

a[a[:, 0].sort()[1]]

Output:

tensor([[3., 5.],
        [3., 7.],
        [5., 5.],
        [5., 3.],
        [6., 4.]])

如果你真的需要它交错：

b = a[a[:, 1].sort()[1]]
b[b[:, 0].sort()[1]]

Output:

tensor([[3., 5.],
        [3., 7.],
        [5., 3.],
        [5., 5.],
        [6., 4.]])