From d9b7af483de2acd3990233631d87211f50859a7a Mon Sep 17 00:00:00 2001
From: Vincent Roseberry <rosbo@google.com>
Date: Fri, 8 Jan 2021 18:02:31 +0000
Subject: [PATCH] Use TensorFlow 2.4.0 wheels.

BUG=172352589
---
 Dockerfile     | 4 ++--
 gpu.Dockerfile | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index effc65d1..f299c24a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,8 +1,8 @@
 # b/157908450 set to latest once numba 0.49.x fixes performance regression for datashader.
 ARG BASE_TAG=m46
-ARG TENSORFLOW_VERSION=2.3.1
+ARG TENSORFLOW_VERSION=2.4.0
 
-FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37-2 as tensorflow_whl
+FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37 as tensorflow_whl
 FROM gcr.io/deeplearning-platform-release/base-cpu:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh
diff --git a/gpu.Dockerfile b/gpu.Dockerfile
index a9ab9ed8..eb37cfd6 100644
--- a/gpu.Dockerfile
+++ b/gpu.Dockerfile
@@ -1,7 +1,7 @@
 ARG BASE_TAG=staging
 
 FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04 AS nvidia
-FROM gcr.io/kaggle-images/python-tensorflow-whl:2.3.1-py37-2 as tensorflow_whl
+FROM gcr.io/kaggle-images/python-tensorflow-whl:2.4.0-py37 as tensorflow_whl
 FROM gcr.io/kaggle-images/python:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh