training.ipynb colab: ModuleNotFoundError: No module named 'brax'

[samuela]

I'm seeing the following error running https://colab.research.google.com/github/google/brax/blob/main/notebooks/training.ipynb in Colab without modification:

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-2-6058247c80ff>](https://localhost:8080/#) in <cell line: 0>()
     14 try:
---> 15   import brax
     16 except ImportError:

ModuleNotFoundError: No module named 'brax'

During handling of the above exception, another exception occurred:

ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-2-6058247c80ff>](https://localhost:8080/#) in <cell line: 0>()
     17   get_ipython().system('pip install git+https://github.com/google/brax.git@main')
     18   clear_output()
---> 19   import brax
     20 
     21 import flax

ModuleNotFoundError: No module named 'brax'

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
---------------------------------------------------------------------------

[btaba]

Looks like ImportError should be switched to ModuleNotFoundError, thanks for the bug report!