pytorch nightly build should be installed, see Start Locally.

git clone https://github.com/xadupre/experimental-experiment.git
pip install onnxruntime-gpu pynvml
pip install -r requirements-dev.txt    
export PYTHONPATH=$PYTHONPATH:<this folder>

Then install onnx-rewriter.

The script evaluates the memory peak, the computation time of the exporters. It also compares the exported models when run through onnxruntime. The full script takes around 20 minutes to complete. It stores on disk all the graphs, the data used to draw them, and the models.

python _doc/examples/plot_torch_export.py -s large