Commit bc42fea
| 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" | ||
0 commit comments