[MPS] Fix error check for torch.var on scalar

malfet · malfet · commit e691c86ed6fe · 2025-08-18T08:30:40.000-07:00
Fixes #160738 ghstack-source-id: 8102c54 Pull Request resolved: #160889
diff --git a/aten/src/ATen/native/mps/operations/ReduceOps.mm b/aten/src/ATen/native/mps/operations/ReduceOps.mm
@@ -456,7 +456,7 @@ static Tensor std_var_common_impl_mps(const Tensor& input_t,
     errMessage += ": reduction dim must be in the range of input shape";
     for (const auto dim : dim_value) {
       auto wrap_dim = maybe_wrap_dim(dim, num_input_dims);
-      TORCH_CHECK(wrap_dim < static_cast<decltype(wrap_dim)>(input_shape.size()), errMessage.c_str())
+      TORCH_CHECK(wrap_dim < (num_input_dims ? num_input_dims : 1), errMessage.c_str())
     }
   }
 
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -5321,6 +5321,9 @@ def helper():
 
         helper()
 
+        # Regression test for https://github.com/pytorch/pytorch/issues/160738
+        self.assertTrue(torch.var(torch.tensor(3.13, device='mps'), dim=0).isnan().item())
+
     # Test forward amax
     def test_amax(self):
         def helper(shape, dim, keepdim):

Original file line number	Diff line number	Diff line change
`@@ -456,7 +456,7 @@ static Tensor std_var_common_impl_mps(const Tensor& input_t,`
`456`	`456`	`errMessage += ": reduction dim must be in the range of input shape";`
`457`	`457`	`for (const auto dim : dim_value) {`
`458`	`458`	`auto wrap_dim = maybe_wrap_dim(dim, num_input_dims);`
`459`		`- TORCH_CHECK(wrap_dim < static_cast<decltype(wrap_dim)>(input_shape.size()), errMessage.c_str())`
	`459`	`+ TORCH_CHECK(wrap_dim < (num_input_dims ? num_input_dims : 1), errMessage.c_str())`
`460`	`460`	`}`
`461`	`461`	`}`
`462`	`462`