Update

malfet · malfet · commit a672811d60f4 · 2025-04-21T16:07:01.000-07:00
[ghstack-poisoned]
diff --git a/aten/src/ATen/native/mkldnn/Utils.cpp b/aten/src/ATen/native/mkldnn/Utils.cpp
@@ -19,6 +19,7 @@ std::vector<int64_t> pool_output_sizes(
   output_size[1] = input_size[1];
 
   for (const auto i : c10::irange(2, input_size.size())) {
+    TORCH_CHECK_VALUE(stride[i -2] > 0, "Strides must be positive!");
     output_size[i] = pooling_output_shape_pad_lr<int64_t>(
       input_size[i],
       kernel_size[i - 2],
diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -1623,6 +1623,12 @@ def test_mkldnn_setflags_nowarn(self, device):
         # Above should trigger no warnings regardless of configuration
         self.assertEqual(len(w), 0)
 
+    def test_mkldnn_error_on_zero_stride(self, device):
+        # Regression test for https://github.com/pytorch/pytorch/issues/149274
+        x = torch.rand(1, 2, 3, 3).to_mkldnn()
+        with self.assertRaises(ValueError):
+            torch.mkldnn_max_pool2d(x, kernel_size=3, stride=0)
+
 
 instantiate_device_type_tests(TestMkldnn, globals(), only_for=('cpu',))
 
