1import torch 2from executorch.exir import to_edge 3from torch.export import export 4 5 6# Start with a PyTorch model that adds two input tensors (matrices) 7class Add(torch.nn.Module): 8 def __init__(self): 9 super(Add, self).__init__() 10 11 def forward(self, x: torch.Tensor, y: torch.Tensor): 12 return x + y 13 14 15# 1. torch.export: Defines the program with the ATen operator set. 16aten_dialect = export(Add(), (torch.ones(1), torch.ones(1))) 17 18# 2. to_edge: Make optimizations for Edge devices 19edge_program = to_edge(aten_dialect) 20 21# 3. to_executorch: Convert the graph to an ExecuTorch program 22executorch_program = edge_program.to_executorch() 23 24# 4. Save the compiled .pte program 25with open("add.pte", "wb") as file: 26 file.write(executorch_program.buffer) 27