justinchuby / torch-onnx

A standalone version of the next PyTorch ONNX exporter
MIT License
2 stars 1 forks source link

Support input_names in dynamic_axes #85

Closed titaiwangms closed 3 months ago

titaiwangms commented 3 months ago

Fix https://github.com/pytorch/pytorch/issues/129276

Tested with

    def test_input_names_are_not_yet_supported_in_dynamic_axes(self):
        import torch_onnx
        torch_onnx.patch_torch(error_report=True, profile=True)

        a = torch.onnx.export(
            SampleModelForDynamicShapes(),
            (
                torch.randn(2, 2, 3),
                torch.randn(2, 2, 3),
            ),
            input_names=["input"],
            dynamic_axes={"input": [0, 1]},
            dynamo=True,
        )
        a.save("test.onnx")

Screenshot 2024-07-05 111822