skip test_sort_large on MI350

Author: dnikolaev-amd

aten/src/ATen/native/cuda/SortStable.cu

@@ -230,6 +230,13 @@ void launch_stable_sort_kernel(
       std::min(numel, static_cast<int64_t>(std::numeric_limits<int>::max()));
   int64_t nsort = self.size(dim);
   int64_t nbatch = (numel_or_intmax / nsort) * nsort;
+  TORCH_WARN("\n^^^^^^^^^^^^^^^^^^^^^^^^^",
+            "\ndim=", dim,
+            "\nnumel=", numel,
+            "\nnumel_or_intmax=", numel_or_intmax,
+            "\nnsort=", nsort,
+            "\nnbatch=", nbatch,
+  )
   TORCH_CHECK(nbatch > 0, "Cannot sort dimension of length ", nsort);
   int64_t* indices_ptr = indices.mutable_data_ptr<int64_t>();
 

test/test_sort_and_select.py

@@ -26,6 +26,7 @@
 )
 from torch.testing._internal.common_utils import (
     run_tests,
+    skipIfRocmArch,
     skipIfTorchDynamo,
     slowTest,
     TestCase,
@@ -208,6 +209,7 @@ def test_stable_sort(self, device, dtype):
                 torch.arange(start=1, end=2 * ncopies, step=2, device=device),
             )
 
+    @skipIfRocmArch(("gfx950", ))
     @onlyCUDA
     @dtypes(torch.uint8)
     @largeTensorTest("200GB")  # Unfortunately 80GB A100 is not large enough