For what it’s worth, a slightly more concise way to do this (that also avoids creating a new tensor, which I find can be cumbersome for autograd purposes) would be
>>> x = torch.tensor([
        [1, 0, 9, 87, 1],
        [6, 0, 8, 1, 8],
        [4, 0, 7, 2, 2],
        [3, 0, 2, 6, 3]
    ])
>>> _, indices = torch.topk(x[:, 0], 2)
>>> indices
tensor([1, 2])
>>> new_dataset = x[indices, 2:5]
>>> new_dataset
tensor([[8., 1., 8.],
        [7., 2., 2.]])