Polygen: model_test.py not working

System details:
Ubuntu 20.04 with python 3.7.10

Followed installation instructions in README and ran run.sh

I've got the following error when the following command executed - `python3 model_test.py`
```
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint8 = np.dtype([("qint8", np.int8, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint8 = np.dtype([("quint8", np.uint8, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint16 = np.dtype([("qint16", np.int16, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint16 = np.dtype([("quint16", np.uint16, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint32 = np.dtype([("qint32", np.int32, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
np_resource = np.dtype([("resource", np.ubyte, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint8 = np.dtype([("qint8", np.int8, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint8 = np.dtype([("quint8", np.uint8, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint16 = np.dtype([("qint16", np.int16, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint16 = np.dtype([("quint16", np.uint16, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint32 = np.dtype([("qint32", np.int32, 1)])
/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
np_resource = np.dtype([("resource", np.ubyte, 1)])
Traceback (most recent call last):
File "/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/sonnet/__init__.py", line 47, in _ensure_dependency_available_at_version
pkg = importlib.import_module(package_name)
File "/home/mppavan/.pyenv/versions/3.7.10/lib/python3.7/importlib/__init__.py", line 127, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "", line 1006, in _gcd_import
File "", line 983, in _find_and_load
File "", line 967, in _find_and_load_unlocked
File "", line 677, in _load_unlocked
File "", line 728, in exec_module
File "", line 219, in _call_with_frames_removed
File "/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow_probability/__init__.py", line 68, in
_ensure_tf_install()
File "/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/tensorflow_probability/__init__.py", line 65, in _ensure_tf_install
present=tf.__version__))
ImportError: This version of TensorFlow Probability requires TensorFlow version >= 1.15; Detected an installation of version 1.14.0. Please upgrade TensorFlow to proceed.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "model_test.py", line 16, in
from modules import FaceModel
File "/home/mppavan/deepmind/deepmind-research/polygen/modules.py", line 16, in
import sonnet as snt
File "/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/sonnet/__init__.py", line 64, in
_ensure_dependency_available_at_version('tensorflow_probability', '0.4.0')
File "/home/mppavan/deepmind/deepmind-research/polygen/polygen/lib/python3.7/site-packages/sonnet/__init__.py", line 54, in _ensure_dependency_available_at_version
package_name, min_version, pip_name, pip_name))
SystemError: Sonnet requires tensorflow_probability (minimum version 0.4.0) to be installed. If using pip, run `pip install tensorflow-probability` or `pip install tensorflow-probability-gpu`
```