File tree Expand file tree Collapse file tree 1 file changed +3
-2
lines changed Expand file tree Collapse file tree 1 file changed +3
-2
lines changed Original file line number Diff line number Diff line change @@ -89,7 +89,8 @@ ENV JAX_PLATFORM=linux_x86_64
89
89
ENV JAX_BASE_URL="https://storage.googleapis.com/jax-releases"
90
90
91
91
RUN pip install $JAX_BASE_URL/$JAX_CUDA_VERSION/jaxlib-$JAXLIB_VERSION-$JAX_PYTHON_VERSION-none-$JAX_PLATFORM.whl && \
92
- pip install jax==$JAX_VERSION
92
+ pip install jax==$JAX_VERSION && \
93
+ /tmp/clean-layer.sh
93
94
94
95
# Reinstall packages with a separate version for GPU support.
95
96
COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/
@@ -113,4 +114,4 @@ RUN pip install pycuda && \
113
114
# ADD patches/tensorboard/notebook.py /opt/conda/lib/python3.7/site-packages/tensorboard/notebook.py
114
115
115
116
# Remove the CUDA stubs.
116
- ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH_NO_STUBS"
117
+ ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH_NO_STUBS"
You can’t perform that action at this time.
0 commit comments