diff --git a/.travis.yml b/.travis.yml
index 4995d8d..1107fc5 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -26,6 +26,7 @@ env:
     - PROJECT="transporter"
     - PROJECT="tvt"
     # - PROJECT="unrestricted_advx"  # TODO(b/184862249): Fix and enable
+    - PROJECT="unsupervised_adversarial_training"
 before_script:
   - sudo apt-get update -qq
   - pip install --upgrade pip
diff --git a/unsupervised_adversarial_training/quick_eval_cifar.py b/unsupervised_adversarial_training/quick_eval_cifar.py
index 249a628..9dac772 100644
--- a/unsupervised_adversarial_training/quick_eval_cifar.py
+++ b/unsupervised_adversarial_training/quick_eval_cifar.py
@@ -30,7 +30,6 @@ import cleverhans
 from cleverhans import attacks
 import numpy as np
 import tensorflow.compat.v1 as tf
-from tensorflow.python.ops import math_grad
 import tensorflow_hub as hub
 
 UAT_HUB_URL = ('https://tfhub.dev/deepmind/unsupervised-adversarial-training/'
@@ -80,10 +79,6 @@ def eval_cifar():
   _, data_test = tf.keras.datasets.cifar10.load_data()
   data = _build_dataset(data_test, batch_size=batch_size, shuffle=False)
 
-  # Necessary for backwards-compatibility
-  # Earlier versions of TF don't have a registered gradient for the AddV2 op
-  tf.RegisterGradient('AddV2')(math_grad._AddGrad)  # pylint: disable=protected-access
-
   # Generate adversarial images.
   if attack_fn_name == 'fgsm':
     attack = attacks.MadryEtAl(cleverhans_model)
diff --git a/unsupervised_adversarial_training/requirements.txt b/unsupervised_adversarial_training/requirements.txt
index d49ab6d..7e95236 100644
--- a/unsupervised_adversarial_training/requirements.txt
+++ b/unsupervised_adversarial_training/requirements.txt
@@ -1,7 +1,7 @@
 absl-py>=0.7.0
-cleverhans>=3.0.1
+cleverhans>=3.0.1,<=3.1.0
 numpy>=1.16.4
 pillow>=4.3.0
-tensorflow>=1.14
-tensorflow-gpu  >= 1.11.0  # GPU version of TensorFlow.
+tensorflow>=1.15,<2
+# tensorflow-gpu >= 1.15.0,<2  # GPU version of TensorFlow.
 tensorflow-hub>=0.5.0
diff --git a/unsupervised_adversarial_training/run.sh b/unsupervised_adversarial_training/run.sh
index c239c4e..0d15ff7 100755
--- a/unsupervised_adversarial_training/run.sh
+++ b/unsupervised_adversarial_training/run.sh
@@ -17,9 +17,11 @@
 # user@host:/path/to/deepmind_research$ unsupervised_adversarial_training/run.sh
 
 # Sets up virtual environment, install dependencies, and runs evaluation script
-python3 -m venv uat_venv
-source uat_venv/bin/activate
+python3 -m venv /tmp/uat_venv
+source /tmp/uat_venv/bin/activate
+pip install -U pip
 pip install -r unsupervised_adversarial_training/requirements.txt
 
 python -m unsupervised_adversarial_training.quick_eval_cifar \
-  --attack_fn_name=fgsm
+  --attack_fn_name=fgsm \
+  --num_batches=1