diff --git a/README.md b/README.md
index 2285001a..b9ec9a63 100644
--- a/README.md
+++ b/README.md
@@ -142,7 +142,7 @@ Please use the following command to create a TPU device in the Cloud VM.
 
 ```sh
 ctpu up --name=$TPU_NAME --project=$PROJECT --zone=$ZONE --tpu-size=v3-8  \
-        --tpu-only   --tf-version=1.15 --noconf
+        --tpu-only   --tf-version=2.1 --noconf
 ```
 
 
diff --git a/setup.py b/setup.py
index f1a234c5..4836ba11 100644
--- a/setup.py
+++ b/setup.py
@@ -56,11 +56,10 @@
         'scipy',
         'sentencepiece',
         'six>=1.14',  # TODO(adarob): Remove once rouge-score is updated.
-        'tensorflow-text~=1.15',
+        'tensorflow-text',
         'tfds-nightly',
     ],
     extras_require={
-        'tensorflow': ['tensorflow~=1.15'],
         'gcp': ['gevent', 'google-api-python-client', 'google-compute-engine',
                 'google-cloud-storage', 'oauth2client'],
         'cache-tasks': [
diff --git a/t5/version.py b/t5/version.py
index 80e40b9b..0609aaf2 100644
--- a/t5/version.py
+++ b/t5/version.py
@@ -19,4 +19,4 @@
 pulling in all the dependencies in __init__.py.
 """
 
-__version__ = '0.3.2'
+__version__ = '0.4.0'