diff --git a/.travis.yml b/.travis.yml
index 3a91af2da..451504da6 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -3,10 +3,6 @@ python:
   - "3.5"
 before_install:
   - pip install -e .[tests]
-  # Replace tensorflow-gpu by tensorflow, see https://github.com/tensorflow/addons/issues/532
-  - export TF_VERSION="$(pip show tensorflow-gpu | grep Version | cut -f2 -d' ')"
-  - pip uninstall -y tensorflow-gpu
-  - pip install tensorflow==$TF_VERSION
 before_script:
   - wget https://s3.amazonaws.com/opennmt-models/transliteration-aren-v2.tar.gz
   - mkdir -p $TRAVIS_BUILD_DIR/testdata
diff --git a/setup.py b/setup.py
index c0a4f0491..fadfad97a 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@
         "pyyaml==5.1.*",
         "rouge==0.3.1",
         "sacrebleu>=1.3.6,<2;python_version>='3.0'",
-        "tensorflow-addons>=0.5.1"
+        "tensorflow-addons>=0.5.2"
     ],
     extras_require={
         "tests": tests_require,