diff --git a/sonnet/src/recurrent.py b/sonnet/src/recurrent.py
index 2c5364d..97b75b7 100644
--- a/sonnet/src/recurrent.py
+++ b/sonnet/src/recurrent.py
@@ -33,7 +33,7 @@
 
 # pylint: disable=g-direct-tensorflow-import
 # Required for specializing `UnrolledLSTM` per device.
-from tensorflow.python import context as context_lib
+from tensorflow.python.eager import context as context_lib
 # pylint: enable=g-direct-tensorflow-import