[MKLDNN] Check that strides are positive (#151848)

malfet · pytorchbot · commit 8d692968ae48 · 2025-05-07T19:59:19.000Z
For pooling ops. Prevents division-by-zero when argument is wrong Fixes #149274 Pull Request resolved: #151848 Approved by: https://github.com/atalman (cherry picked from commit 6f32712)
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',))
 
