Add a way to extract an `ONNX` file

[Luni-4]

First of all, thanks a lot for this amazing project and your hard work!

In https://github.com/burn-rs/burn we are trying to load different ONNX models, convert them into burn models, and then use them on different backends. It would be fantastic if we could try your model too. Can you add a way to extract the ONNX model? Using PyTorch is quite feasible.

Thanks a lot in advance!