Fix when _stable_1d_sort to work when n >= N (#6177)

* Fix when _stable_1d_sort to work when n >= N

* Apply suggestions

Co-authored-by: Carlos Mocholi <[email protected]>

Author: frankier

pytorch_lightning/metrics/utils.py

@@ -285,6 +285,7 @@ def _stable_1d_sort(x: torch, N: int = 2049):
     n = x.numel()
     if N - n > 0:
         x_max = x.max()
-        x_pad = torch.cat([x, (x_max + 1) * torch.ones(2049 - n, dtype=x.dtype, device=x.device)], 0)
-    x_sort = x_pad.sort()
-    return x_sort.values[:n], x_sort.indices[:n]
+        x = torch.cat([x, (x_max + 1) * torch.ones(N - n, dtype=x.dtype, device=x.device)], 0)
+    x_sort = x.sort()
+    i = min(N, n)
+    return x_sort.values[:i], x_sort.indices[:i]

tests/metrics/classification/test_auc.py

@@ -61,4 +61,4 @@ def test_auc_functional(self, x, y):
 ])
 def test_auc(x, y, expected):
     # Test Area Under Curve (AUC) computation
-    assert auc(torch.tensor(x), torch.tensor(y)) == expected
+    assert auc(torch.tensor(x), torch.tensor(y), reorder=True) == expected