import torch
import torch.nn as nn

class AddModel(nn.Module):
    def forward(self, x1, x2):
        return x1 + x2

model = AddModel()
x1 = torch.randn(1, 3)
x2 = torch.randn(1, 3)

torch.onnx.export(
    model,
    (x1, x2),
    "add.onnx",
    input_names=["x1", "x2"],
    output_names=["y"],
    dynamic_axes={
        "x1": {0: "batch"},
        "x2": {0: "batch"},
        "y": {0: "batch"},
    },
    opset_version=18,
)