diff --git a/edward/util/random_variables.py b/edward/util/random_variables.py
index 3a581505a..a706645f1 100644
--- a/edward/util/random_variables.py
+++ b/edward/util/random_variables.py
@@ -12,7 +12,7 @@
 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.framework.ops import set_shape_and_handle_data_for_outputs as set_shapes_for_outputs
 from tensorflow.python.util import compat
 
 tfb = tf.contrib.distributions.bijectors