From 2124b24a1ad76aac6f4595a1ac834715dffe1230 Mon Sep 17 00:00:00 2001 From: Toby Roseman Date: Fri, 10 Jul 2026 14:05:20 -0700 Subject: [PATCH] Bug fix: max_pool2d uses default stride value --- coreai_torch/_aten_to_core.py | 4 ++-- tests/ops/test_ops.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index cd342fa..3a38760 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -2227,11 +2227,11 @@ def replace_maxpool2d_with_indices( x = _get_operand(values_map, node, 0) args = node.args kernel_size = args[1] - if isinstance(args[2], fx.Node): + if len(args) > 2 and isinstance(args[2], fx.Node): raise ValueError( f"Encountered dynamic stride at maxpool2d: node: {node}, name: {node.name}" ) - stride = args[2] + stride = args[2] if len(args) >= 3 else kernel_size padding = args[3] if len(args) >= 4 else [0, 0] dilation = args[4] if len(args) >= 5 else [1, 1] ceil_mode = args[5] if len(args) >= 6 else False diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index a19fb91..cb5bcb5 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -2686,6 +2686,8 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor: @pytest.mark.parametrize( "input_shape, dtype, kernel_size, stride, padding, dilation, ceil_mode, dynamic_dims", [ + # Default parameters + ((2, 4, 16, 16), torch.float32, 3, None, 0, 1, False, tuple()), # Static — all pool configs, multiple shapes and dtypes ((2, 4, 16, 16), torch.float32, 3, 2, 1, 1, False, tuple()), ((2, 4, 16, 16), torch.float32, 3, 2, 0, 1, True, tuple()),