Use TensorFlow 2.4.0 wheels. (Kaggle#931)

BUG=172352589
erdes · Jan 11, 2021 · 67cf125 · 67cf125
1 parent b9bd8f0
commit 67cf125
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git 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
@@ -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