RTDataScience
diff --git a/‎Dockerfile‎
Lines changed: 4 additions & 1 deletion b/‎Dockerfile‎
Lines changed: 4 additions & 1 deletion
diff --git a/‎gpu.Dockerfile‎
Lines changed: 3 additions & 1 deletion b/‎gpu.Dockerfile‎
Lines changed: 3 additions & 1 deletion
diff --git a/‎tests/test_tensorflow.py‎
Lines changed: 1 addition & 1 deletion b/‎tests/test_tensorflow.py‎
Lines changed: 1 addition & 1 deletion
@@ -51,7 +51,10 @@ RUN pip install seaborn python-dateutil dask && \
     make -j $(nproc) && make install && \
     /tmp/clean-layer.sh
 
-RUN pip install tensorflow==2.0.0rc1 && \
+# Install tensorflow from a pre-built wheel
+COPY --from=tensorflow_whl /tmp/tensorflow_cpu/*.whl /tmp/tensorflow_cpu/
+RUN pip install /tmp/tensorflow_cpu/tensorflow*.whl && \
+    rm -rf /tmp/tensorflow_cpu && \
     /tmp/clean-layer.sh
 
 RUN apt-get install -y libfreetype6-dev && \
 
@@ -1,6 +1,7 @@
 ARG BASE_TAG=staging
 
 FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu16.04 AS nvidia
+FROM gcr.io/kaggle-images/python-tensorflow-whl:2.0.0-rc1-py36 as tensorflow_whl
 FROM gcr.io/kaggle-images/python:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh
@@ -44,8 +45,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     /tmp/clean-layer.sh
 
 # Reinstall packages with a separate version for GPU support.
+COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/
 RUN pip uninstall -y tensorflow && \
-    pip install tensorflow-gpu==2.0.0rc1 && \
+    pip install /tmp/tensorflow_gpu/tensorflow*.whl && \
     rm -rf /tmp/tensorflow_gpu && \
     conda remove --force -y pytorch-cpu torchvision-cpu && \
     conda install -y pytorch torchvision cudatoolkit=10.0 -c pytorch && \
 
@@ -8,7 +8,7 @@
 
 class TestTensorflow(unittest.TestCase):
     def test_addition(self):        
-        result = tf.add(2, 3)        
+        result = tf.add(2, 3)
         self.assertEqual(5, result.numpy())
 
     def test_conv2d(self):