fix aminmax output resize issue when input is a zero dimension tensor

ghstack-source-id: 9689098389a896779c3fc31d8605971516667437 Pull Request resolved: #96171
pytorch · Mar 7, 2023 · cae5db1 · cae5db1
1 parent cc775fb
commit cae5db1
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp
@@ -177,6 +177,14 @@ static void aminmax_kernel(
     "Expect min and max dtype ", self.scalar_type(),
     " but got ", min_result.scalar_type(), " and ", max_result.scalar_type());
 
+  if (self.numel() == 1 && self.ndimension() == 0) {
+    min_result.resize_({});
+    max_result.resize_({});
+    min_result.fill_(self);
+    max_result.fill_(self);
+    return;
+  }
+
   AT_DISPATCH_ALL_TYPES_AND(ScalarType::Bool, self.scalar_type(), "aminmax_cpu", [&] {
     compare_base_kernel<scalar_t, scalar_t>(min_result, max_result, self, wrap_dim, keepdim, [&] (
       scalar_t* min_result_data, scalar_t* max_result_data,