diff --git a/examples/arm/example_modules/README.md b/examples/arm/example_modules/README.md new file mode 100644 index 00000000000..9a746114b98 --- /dev/null +++ b/examples/arm/example_modules/README.md @@ -0,0 +1,7 @@ +# Example of an external model for the ARM AOT Compiler +Example of an external Python file to be used as a module by the `run.sh` (and the `aot_arm_compiler.py`) scripts in `examples/arm` directory. +Just pass the path of the `add.py` file as `--model_name`: + +`ModelUnderTest` should be a `torch.nn.module` instance. + +`ModelInputs` should be a tuple of inputs to the forward function. diff --git a/examples/arm/example_modules/add.py b/examples/arm/example_modules/add.py new file mode 100644 index 00000000000..6942e97f807 --- /dev/null +++ b/examples/arm/example_modules/add.py @@ -0,0 +1,13 @@ +import torch + + +class myModelAdd(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x + x + + +ModelUnderTest = myModelAdd() +ModelInputs = (torch.ones(5),)