diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 44cfd51a3..010fb6736 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -70,14 +70,10 @@ def _sort_single_gradient_block(
             _dispatch.to_index_array(sample_values[:, 0])
         ]
 
-        # I AM NOT SURE WHAT THESE LINES DO THE OLD TEST PASS EVEN
-        # IF THEY ARE COMMENTED
-
         # sort the samples in gradient regularly moving the rows considering all columns
-        # sorted_idx = _dispatch.argsort_labels_values(sample_values,
-        # reverse=descending)
-        # sample_values = sample_values[sorted_idx]
-        # values = values[sorted_idx]
+        sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
+        sample_values = sample_values[sorted_idx]
+        values = values[sorted_idx]
     if "components" in axes:
         for i, _ in enumerate(gradient_block.components):
             sorted_idx = _dispatch.argsort_labels_values(
@@ -254,24 +250,6 @@ def sort_block(
     >>> sorted_block = metatensor.sort_block(block)
     >>> np.all(sorted_block.values == block_sorted_stepwise.values)
     True
-    >>> # You can also choose along which axis of "samples“ you sort
-    >>> block2 = TensorBlock(
-    ...     values=np.arange(12).reshape(4, 3),
-    ...     samples=Labels(
-    ...         ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
-    ...     ),
-    ...     components=[],
-    ...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
-    ... )
-    >>> block_sorted_2_sample = metatensor.sort_block(
-    ...     block2, axes=["samples"], name="atom"
-    ... )
-    >>> # samples (first dimension of the array) are sorted
-    >>> block_sorted_2_sample.values
-    array([[ 3,  4,  5],
-           [ 9, 10, 11],
-           [ 0,  1,  2],
-           [ 6,  7,  8]])
     >>> # This function can also sort gradients:
     >>> sorted_block.add_gradient(
     ...     parameter="g",
@@ -297,7 +275,24 @@ def sort_block(
     <BLANKLINE>
            [[ 6,  7,  8],
             [ 9, 10, 11]]])
-
+    >>> # You can also choose along which axis of "samples“ you sort
+    >>> block2 = TensorBlock(
+    ...     values=np.arange(12).reshape(4, 3),
+    ...     samples=Labels(
+    ...         ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
+    ...     ),
+    ...     components=[],
+    ...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
+    ... )
+    >>> block_sorted_2_sample = metatensor.sort_block(
+    ...     block2, axes=["samples"], name="atom"
+    ... )
+    >>> # samples (first dimension of the array) are sorted
+    >>> block_sorted_2_sample.values
+    array([[ 3,  4,  5],
+           [ 9, 10, 11],
+           [ 0,  1,  2],
+           [ 6,  7,  8]])
     """
     if isinstance(axes, str):
         if axes == "all":
diff --git a/python/metatensor-operations/tests/sort.py b/python/metatensor-operations/tests/sort.py
index 9c8138d32..a15c57b7d 100644
--- a/python/metatensor-operations/tests/sort.py
+++ b/python/metatensor-operations/tests/sort.py
@@ -341,26 +341,26 @@ def tensor_two_samples_ascending_a():
         gradient=TensorBlock(
             values=np.array(
                 [
-                    [[1, 2], [3, 4]],
                     [[5, 6], [7, 8]],
+                    [[17, 18], [19, 20]],
                     [[9, 10], [11, 12]],
                     [[13, 14], [15, 16]],
-                    [[17, 18], [19, 20]],
-                    [[21, 22], [23, 24]],
+                    [[1, 2], [3, 4]],
                     [[25, 26], [27, 28]],
+                    [[21, 22], [23, 24]],
                 ]
             ),
             samples=Labels(
                 ["sample", "g"],
                 np.array(
                     [
-                        [3, 1],
                         [0, 1],
+                        [1, 1],
                         [2, 1],
                         [2, 2],
-                        [1, 1],
-                        [4, 1],
+                        [3, 1],
                         [4, 0],
+                        [4, 1],
                     ]
                 ),
             ),
@@ -451,6 +451,14 @@ def test_raise_error(tensor, tensor_sorted_ascending):
 
 
 def test_sort_two_sample(tensor_two_samples, tensor_two_samples_ascending_a):
+    print(
+        "ii",
+        metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a")
+        .gradient("g")
+        .samples,
+    )
+    print("jj", tensor_two_samples_ascending_a.block(0).gradient("g").samples)
+
     metatensor.allclose_block_raise(
         metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a"),
         tensor_two_samples_ascending_a.block(0),