Flax GPT-2 HuggingFace Model

In this example, we show how to optimize Flax GPT-2 model from the HuggingFace repository. We recommend running this example in NVIDIA NGC TensorFlow2 container.

Install dependencies and set up the environment

. ./install.sh

Run optimization

./optimize.py