File tree Expand file tree Collapse file tree 1 file changed +5
-3
lines changed
aten/src/ATen/native/cuda Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Original file line number Diff line number Diff line change @@ -225,8 +225,9 @@ void launch_stable_sort_kernel(
225225 return ;
226226 }
227227
228- int64_t numel_or_intmax =
229- std::min (numel, static_cast <int64_t >(std::numeric_limits<int >::max ()));
228+ const int64_t intmax = static_cast <int64_t >(std::numeric_limits<int >::max ());
229+ // On ROCm, std::min -> ::min did not work as expected on when input values >= 2147483648
230+ int64_t numel_or_intmax = numel < intmax ? numel : intmax;
230231 int64_t nsort = self.size (dim);
231232 int64_t nbatch = (numel_or_intmax / nsort) * nsort;
232233 TORCH_CHECK (nbatch > 0 , " Cannot sort dimension of length " , nsort);
@@ -238,7 +239,8 @@ void launch_stable_sort_kernel(
238239 scalar_t * values_ptr = values.mutable_data_ptr <scalar_t >();
239240 int64_t remaining = numel;
240241 while (remaining > 0 ) {
241- int64_t n = std::min (remaining, nbatch);
242+ // On ROCm, std::min -> ::min did not work as expected on when input values >= 2147483648
243+ int64_t n = remaining < nbatch ? remaining : nbatch;
242244 int64_t nsegments = n / nsort;
243245
244246 if (nsegments == 1 ||
You can’t perform that action at this time.
0 commit comments