With Flax, you can profile your code by adding the --profile_steps==5 parameter to your training command. Install the Tensorboard profile plugin:   Copied pip install tensorflow tensorboard-plugin-profile