| 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, | |
| ) |