Merge pull request #809 from Kaggle/clean-layer-jax

Call clean-layer script after JAX installation

Author: rosbo

gpu.Dockerfile

@@ -89,7 +89,8 @@ ENV JAX_PLATFORM=linux_x86_64
 ENV JAX_BASE_URL="https://storage.googleapis.com/jax-releases"
 
 RUN  pip install $JAX_BASE_URL/$JAX_CUDA_VERSION/jaxlib-$JAXLIB_VERSION-$JAX_PYTHON_VERSION-none-$JAX_PLATFORM.whl && \
-     pip install jax==$JAX_VERSION
+     pip install jax==$JAX_VERSION && \
+     /tmp/clean-layer.sh
 
 # Reinstall packages with a separate version for GPU support.
 COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/
@@ -113,4 +114,4 @@ RUN pip install pycuda && \
 # ADD patches/tensorboard/notebook.py /opt/conda/lib/python3.7/site-packages/tensorboard/notebook.py
 
 # Remove the CUDA stubs.
-ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH_NO_STUBS"
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH_NO_STUBS"