HelloWorld2307's picture
Upload ops folder as ops_data
bddb894 verified
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,
)