From 85b9f0fd3d831d9789bfa8038510b0433095c6fb Mon Sep 17 00:00:00 2001
From: Reuben Morais <reuben.morais@gmail.com>
Date: Sat, 2 Jan 2021 13:28:04 +0000
Subject: [PATCH] Bump training dependencies to TensorFlow 2.3.1

---
 Dockerfile.train.tmpl | 2 +-
 doc/TRAINING.rst      | 6 +++---
 setup.py              | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Dockerfile.train.tmpl b/Dockerfile.train.tmpl
index 9baa76d9..1b49513a 100644
--- a/Dockerfile.train.tmpl
+++ b/Dockerfile.train.tmpl
@@ -1,6 +1,6 @@
 # Please refer to the TRAINING documentation, "Basic Dockerfile for training"
 
-FROM tensorflow/tensorflow:1.15.4-gpu-py3
+FROM tensorflow/tensorflow:2.3.1-gpu
 ENV DEBIAN_FRONTEND=noninteractive
 
 ENV DEEPSPEECH_REPO=#DEEPSPEECH_REPO#
diff --git a/doc/TRAINING.rst b/doc/TRAINING.rst
index 037dc2b1..07cd75fc 100644
--- a/doc/TRAINING.rst
+++ b/doc/TRAINING.rst
@@ -10,7 +10,7 @@ Prerequisites for training a model
 
 * `Python 3.6 <https://www.python.org/>`_
 * Mac or Linux environment
-* CUDA 10.0 / CuDNN v7.6 per `Dockerfile <https://hub.docker.com/layers/tensorflow/tensorflow/1.15.4-gpu-py3/images/sha256-a5255ae38bcce7c7610816c778244309f8b8d1576e2c0023c685c011392958d7?context=explore>`_.
+* CUDA 10.1 / CuDNN v7.6 per `Dockerfile <https://hub.docker.com/layers/tensorflow/tensorflow/2.3.1-gpu/images/sha256-1d0736e46ae9a961c2111394a43e0bfd266e6151a90d613b6f86229cf01e40e5?context=explore  >`_.
 
 Getting the training code
 ^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -72,7 +72,7 @@ If you have a capable (NVIDIA, at least 8GB of VRAM) GPU, it is highly recommend
 .. code-block:: bash
 
    pip3 uninstall tensorflow
-   pip3 install 'tensorflow-gpu==1.15.4'
+   pip3 install 'tensorflow-gpu==2.3.1'
 
 Please ensure you have the required `CUDA dependency <https://www.tensorflow.org/install/source#gpu>`_ and/or :ref:`Prerequisites <cuda-training-deps>`.
 
@@ -83,7 +83,7 @@ It has been reported for some people failure at training:
    tensorflow.python.framework.errors_impl.UnknownError: Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.
         [[{{node tower_0/conv1d/Conv2D}}]]
 
-Setting the ``TF_FORCE_GPU_ALLOW_GROWTH`` environment variable to ``true`` seems to help in such cases. This could also be due to an incorrect version of libcudnn. Double check your versions with the :ref:`TensorFlow 1.15 documentation <cuda-training-deps>`.
+Setting the ``TF_FORCE_GPU_ALLOW_GROWTH`` environment variable to ``true`` seems to help in such cases. This could also be due to an incorrect version of libcudnn. Double check your versions with the :ref:`TensorFlow 2.3 documentation <cuda-training-deps>`.
 
 Basic Dockerfile for training
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/setup.py b/setup.py
index da705262..63ef8cfe 100644
--- a/setup.py
+++ b/setup.py
@@ -74,7 +74,7 @@ def main():
     ]
 
     tensorflow_pypi_dep = [
-        'tensorflow == 1.15.4'
+        'tensorflow == 2.3.1'
     ]
 
     # Due to pip craziness environment variables are the only consistent way to