input_utils to importing in Colab notebook

In the plex_vit_demo.ipynb notebook demo, the library input_utils is not importing. It gives the errors:

```
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
[](https://3eh0tszcb2z-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20230227-060155-RC00_512581258#) in
16 import uncertainty_baselines as ub
17 import checkpoint_utils # local file import from baselines.jft
---> 18 import input_utils # local file import from baselines.jft
19 import ood_utils # local file import from baselines.jft
20 import preprocess_utils # local file import from baselines.jft

4 frames
[/usr/local/lib/python3.8/dist-packages/uncertainty_baselines/datasets/toxic_comments.py](https://3eh0tszcb2z-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20230227-060155-RC00_512581258#) in _KeyValueStore()
185
186 def lookup(self, keys: tf.Tensor,
--> 187 value_names: Sequence[str]) -> dict[str, tf.Tensor]:
188 """Searchs values of `value_names` by `keys`."""
189 return {

TypeError: 'type' object is not subscriptable
```