bert_dropout_test is failing

[dustinvtran]

bert_dropout_test.py is failing with following error:

  File "/home/runner/work/uncertainty-baselines/uncertainty-baselines/uncertainty_baselines/models/bert_dropout_test.py", line 22, in <module>
    from uncertainty_baselines.models import bert_dropout
  File "/home/runner/work/uncertainty-baselines/uncertainty-baselines/uncertainty_baselines/models/bert_dropout.py", line 24, in <module>
    from official.nlp.modeling import layers as bert_layers
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/official/nlp/modeling/__init__.py", line 21, in <module>
    from official.nlp.modeling import layers
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/official/nlp/modeling/layers/__init__.py", line 47, in <module>
    from official.nlp.modeling.layers.text_layers import BertPackInputs
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/official/nlp/modeling/layers/text_layers.py", line 23, in <module>
    import tensorflow_text as text
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/tensorflow_text/__init__.py", line 21, in <module>
    from tensorflow_text.python import keras
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/tensorflow_text/python/keras/__init__.py", line 21, in <module>
    from tensorflow_text.python.keras.layers import *
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/tensorflow_text/python/keras/layers/__init__.py", line 22, in <module>
    from tensorflow_text.python.keras.layers.tokenization_layers import *
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/tensorflow_text/python/keras/layers/tokenization_layers.py", line 25, in <module>
    from tensorflow_text.python.ops import unicode_script_tokenizer
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/tensorflow_text/python/ops/__init__.py", line 34, in <module>
    from tensorflow_text.python.ops.mst_ops import max_spanning_tree
  File "/opt/hostedtoolcache/Python/3.7.12/x64/lib/python3.7/site-packages/tensorflow_text/python/ops/mst_ops.py", line 33, in <module>
    max_spanning_tree = gen_mst_ops.max_spanning_tree
AttributeError: module '012ff3e36e3c24aefc4a3a7b68a03fedd1e7a7e1' has no attribute 'max_spanning_tree'

https://github.com/google/uncertainty-baselines/runs/5061876922?check_suite_focus=true#step:6:181