Find an ONNX export of the PyTorch model or you can use PyTorch to load a model and then export the model to ONNX. Once you have an ONNX model, then you can use axon_onnx to load into Elixir.

I suggest starting with the following notebook that I contributed to Axon, axon/guides/serialization/onnx_to_axon.livemd at main · elixir-nx/axon · GitHub. If you can get the ONNX to Axon side of the process working for you, then try the same thing on another ONNX model.

If you want to see an example of converting a PyTorch model to ONNX, review the Fast.ai code linked in the above livebook.

** BIG WARNING ** Saving ONNX imported models to an Axon file will probably fail. You can check Sean’s ElixirConf presentation for why.

Another warning: If you want to run a model on your CPU, use Torchx or XLA, you should be fine. Models have GPU memory size requirements. When trying to run a model on a GPU, it needs to fit in the available GPU memory.

Expect a blog post on Axon and ONNX at https://alongtheaxon.com/ about the above onnx_to_axon notebook. It is on my short term TODO.
Good Luck.