input_utils to importing in Colab notebook

[JonoSax]

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)
[<ipython-input-8-e243f97fa52b>](https://3eh0tszcb2z-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20230227-060155-RC00_512581258#) in <module>
     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