bert_dropout_test is failing

`bert_dropout_test.py` is failing with following error:

```python
File "/home/runner/work/uncertainty-baselines/uncertainty-baselines/uncertainty_baselines/models/bert_dropout_test.py", line 22, in
from uncertainty_baselines.models import bert_dropout
File "/home/runner/work/uncertainty-baselines/uncertainty-baselines/uncertainty_baselines/models/bert_dropout.py", line 24, in
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
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
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
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
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
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
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
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
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
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