diff --git a/.travis.yml b/.travis.yml
index 1e5c45eac..4f47f9ea7 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,10 +1,14 @@
 sudo: required
 dist: trusty
 language: python
-matrix:
-    include:
-        - python: 2.7
-        - python: 3.4
+python:
+  - 2.7
+  - 3.4
+  - 3.6
+env:
+  - TF=1.3.0rc0
+  - TF=1.5.0
+  - TF=1.7.0
 notifications:
     email: false
 before_install:
@@ -37,7 +41,7 @@ install:
   - pip install keras
   - pip install matplotlib seaborn scipy
   - pip install networkx==1.9.1 observations sklearn
-  - pip install tensorflow==1.5.0
+  - pip install tensorflow==$TF
   - pip install pystan
   - pip install nbformat nbconvert jupyter_client jupyter
   - python setup.py install
diff --git a/edward/util/random_variables.py b/edward/util/random_variables.py
index 3a581505a..80b04e102 100644
--- a/edward/util/random_variables.py
+++ b/edward/util/random_variables.py
@@ -12,8 +12,12 @@
 from edward.models import PointMass
 from edward.util.graphs import random_variables
 from tensorflow.core.framework import attr_value_pb2
-from tensorflow.python.framework.ops import set_shapes_for_outputs
 from tensorflow.python.util import compat
+try:
+  from tensorflow.python.framework.ops import set_shapes_for_outputs
+except ImportError:
+  from tensorflow.python.framework.ops import \
+      set_shape_and_handle_data_for_outputs as set_shapes_for_outputs
 
 tfb = tf.contrib.distributions.bijectors
 
diff --git a/setup.py b/setup.py
index 45bd5742f..331fcf613 100644
--- a/setup.py
+++ b/setup.py
@@ -15,8 +15,8 @@
     install_requires=['numpy>=1.7',
                       'six>=1.10.0'],
     extras_require={
-        'tensorflow': ['tensorflow>=1.2.0rc0'],
-        'tensorflow with gpu': ['tensorflow-gpu>=1.2.0rc0'],
+        'tensorflow': ['tensorflow >=1.3.0rc0, <=1.7.0'],
+        'tensorflow with gpu': ['tensorflow-gpu >=1.3.0rc0, <=1.7.0'],
         'neural networks': ['keras>=2.0.0', 'prettytensor>=0.7.4'],
         'datasets': ['observations>=0.1.2'],
         'notebooks': ['jupyter>=1.0.0'],
diff --git a/tests/criticisms/metrics_test.py b/tests/criticisms/metrics_test.py
index 452e357e8..621d5b00a 100644
--- a/tests/criticisms/metrics_test.py
+++ b/tests/criticisms/metrics_test.py
@@ -46,23 +46,27 @@ class test_metrics_class(tf.test.TestCase):
 
   def _check_averaging(self, metric, y_true, y_pred):
     n_classes = tf.squeeze(tf.shape(y_true)[-1]).eval()
-    class_scores = [metric(y_true[i], y_pred[i]) for i in range(n_classes)]
+    class_scores = [metric(y_true[:, i], y_pred[:, i])
+                    for i in range(n_classes)]
 
     # No averaging
     no_average = metric(y_true, y_pred, average=None)
     expected_no_average = tf.stack(class_scores)
-    self.assertAllEqual(no_average.eval(), expected_no_average.eval())
+    self.assertAllCloseAccordingToType(
+        no_average.eval(), expected_no_average.eval())
 
     # Macro-averaging
     macro_average = metric(y_true, y_pred, average='macro')
     expected_macro_average = tf.reduce_mean(tf.stack(class_scores))
-    self.assertAllEqual(macro_average.eval(), expected_macro_average.eval())
+    self.assertAllCloseAccordingToType(
+        macro_average.eval(), expected_macro_average.eval())
 
     # Micro-averaging
     micro_average = metric(y_true, y_pred, average='micro')
     expected_micro_average = metric(tf.reshape(y_true, [1, -1]),
                                     tf.reshape(y_pred, [1, -1]))
-    self.assertAllEqual(micro_average.eval(), expected_micro_average.eval())
+    self.assertAllCloseAccordingToType(
+        micro_average.eval(), expected_micro_average.eval())
 
   def test_classification_metrics(self):
     with self.test_session():
@@ -106,7 +110,7 @@ def test_specialized_input_output_metrics(self):
   def test_metrics_with_binary_averaging(self):
     with self.test_session():
       y_true = tf.constant([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0], [3.0, 4.0, 5.0]])
-      y_pred = tf.constant([[2.0, 4.0, 6.0], [4.0, 6.0, 8.0], [6.0, 8.0, 10.0]])
+      y_pred = tf.constant([[3.0, 4.0, 3.0], [6.0, 6.0, 4.0], [9.0, 8.0, 5.0]])
       for metric in all_metrics_with_binary_averaging:
         self._check_averaging(metric, y_true, y_pred)