diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index a0d0a0d880..253637ccd2 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -159,6 +159,21 @@ def forward(self, query, key, value, attn_mask): ) _testing.assert_onnx_program(onnx_program) + def test_dynamic_paddings(self): + class Model(torch.nn.Module): + def forward(self, x): + height = x.size(2) # height is SymInt + x = torch.nn.functional.pad(x, (0, 0, 0, height), mode="replicate") + return x + + onnx_program = torch.onnx.export( + Model(), + (torch.rand(1, 1, 1, 1),), + dynamo=True, + dynamic_shapes=({2: torch.export.Dim("H")},), + ) + _testing.assert_onnx_program(onnx_program) + if __name__ == "__main__": unittest.main()