n id="Flax">
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
tensorboard --logdir runs/fill-circle-100steps-20230411_165612/ Then you can inspect the profile at http://localhost:6006/#profile. If you run into version conflicts with the plugin, try uninstalling and reinstalling all versions 