xref: /aosp_15_r20/external/pytorch/test/mobile/nnc/aot_test_model.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1import torch
2from torch import nn
3
4
5class NeuralNetwork(nn.Module):
6    def forward(self, x):
7        return torch.add(x, 10)
8
9
10model = NeuralNetwork()
11script = torch.jit.script(model)
12torch.jit.save(script, "aot_test_model.pt")
13