Skip to content

Commit 2124b24

Browse files
Toby RosemanTobyRoseman
authored andcommitted
Bug fix: max_pool2d uses default stride value
1 parent 1b3cb3b commit 2124b24

2 files changed

Lines changed: 4 additions & 2 deletions

File tree

coreai_torch/_aten_to_core.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2227,11 +2227,11 @@ def replace_maxpool2d_with_indices(
22272227
x = _get_operand(values_map, node, 0)
22282228
args = node.args
22292229
kernel_size = args[1]
2230-
if isinstance(args[2], fx.Node):
2230+
if len(args) > 2 and isinstance(args[2], fx.Node):
22312231
raise ValueError(
22322232
f"Encountered dynamic stride at maxpool2d: node: {node}, name: {node.name}"
22332233
)
2234-
stride = args[2]
2234+
stride = args[2] if len(args) >= 3 else kernel_size
22352235
padding = args[3] if len(args) >= 4 else [0, 0]
22362236
dilation = args[4] if len(args) >= 5 else [1, 1]
22372237
ceil_mode = args[5] if len(args) >= 6 else False

tests/ops/test_ops.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2686,6 +2686,8 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor:
26862686
@pytest.mark.parametrize(
26872687
"input_shape, dtype, kernel_size, stride, padding, dilation, ceil_mode, dynamic_dims",
26882688
[
2689+
# Default parameters
2690+
((2, 4, 16, 16), torch.float32, 3, None, 0, 1, False, tuple()),
26892691
# Static — all pool configs, multiple shapes and dtypes
26902692
((2, 4, 16, 16), torch.float32, 3, 2, 1, 1, False, tuple()),
26912693
((2, 4, 16, 16), torch.float32, 3, 2, 0, 1, True, tuple()),

0 commit comments

Comments
 (0)