Installation issues in Colab

Issue: Bug
Version: Commit: a7c096532f002a213c25dd69cb8c13ae5c516ded
Command Line output is in colab link below

Hi,
I wanted to run the notebook https://github.com/google/uncertainty-baselines/blob/b3686f75a10b1990c09b8eb589657090b8837d2c/baselines/notebooks/Hyperparameter_Ensembles.ipynb in Colab to check the performance of the hyperparameter ensemble method.

Colab link: https://colab.research.google.com/drive/1mA6LX6P3p2O31TM9bSkqMwoIxlVRUWqh?usp=sharing

Therefore, in colab I ran the following command so that all of the packages were installed properly.
`!pip install "git+https://github.com/google/uncertainty-baselines.git#egg=uncertainty_baselines"`

When I started running the notebook sequentially, in cell 5, I got `ModuleNotFoundError: No module named 'robustness_metrics'`. After a quick search I installed `robustness_metrics` using `pip install "git+https://github.com/google-research/robustness_metrics.git#egg=robustness_metrics"` and tried again. Once again, I got a `ModuleNotFoundError` and installed edward2 with `pip install edward2`. After I handle that, once again I get `ModuleNotFoundError` for seqio which I duly installed.

This time however, the `uncertainty_baselines` module is not loaded in at all and it fails when I do `test_dataset = ub.datasets.get(
DATASET,
split=tfds.Split.TEST).load(batch_size=BATCH_SIZE)` and this seems to be due to dependency issues with `flatbuffer` versions. During the imports I get

`ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
tensorflow 2.6.0 requires flatbuffers~=1.12.0, but you have flatbuffers 2.0 which is incompatible.` and subsequently get

`ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
tf-nightly 2.7.0.dev20210915 requires flatbuffers~=2.0, but you have flatbuffers 1.12 which is incompatible.`

All I want to do is to run the jupyter notebook listed above and experiment with the hyperparameter ensembles technique. Could someone please direct me on how I can run the above jupyter notebook?