File tree Expand file tree Collapse file tree 2 files changed +7
-48
lines changed Expand file tree Collapse file tree 2 files changed +7
-48
lines changed Original file line number Diff line number Diff line change @@ -107,15 +107,6 @@ RUN pip uninstall -y pyarrow && \
107
107
/tmp/clean-layer.sh
108
108
{{ end }}
109
109
110
- # Install implicit
111
- {{ if eq .Accelerator "gpu" }}
112
- RUN mamba install -y implicit implicit-proc=*=gpu && \
113
- /tmp/clean-layer.sh
114
- {{ else }}
115
- RUN mamba install -y implicit && \
116
- /tmp/clean-layer.sh
117
- {{ end}}
118
-
119
110
# Install PyTorch
120
111
{{ if eq .Accelerator "gpu" }}
121
112
COPY --from=torch_whl /tmp/whl/*.whl /tmp/torch/
@@ -172,7 +163,9 @@ RUN pip install spacy && \
172
163
{{ if eq .Accelerator "gpu" }}
173
164
# Install GPU-only packages
174
165
# No specific package for nnabla-ext-cuda 11.x minor versions.
175
- RUN pip install pycuda \
166
+ RUN export PATH=/usr/local/cuda/bin:$PATH && \
167
+ export CUDA_ROOT=/usr/local/cuda && \
168
+ pip install pycuda \
176
169
pynvrtc \
177
170
pynvml && \
178
171
/tmp/clean-layer.sh
@@ -650,7 +643,11 @@ RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tens
650
643
# python -m nb_conda_kernels.install --disable
651
644
652
645
# Force only one libcusolver
646
+ {{ if eq .Accelerator "gpu" }}
653
647
RUN rm /opt/conda/bin/../lib/libcusolver.so.11 && ln -s /usr/local/cuda/lib64/libcusolver.so.11 /opt/conda/bin/../lib/libcusolver.so.11
648
+ {{ else }}
649
+ RUN ln -s /usr/local/cuda/lib64/libcusolver.so.11 /opt/conda/bin/../lib/libcusolver.so.11
650
+ {{ end }}
654
651
655
652
# b/270147159 conda ships with a version of libtinfo which is missing version info causing warnings, replace it with a good version.
656
653
RUN rm /opt/conda/lib/libtinfo.so.6 && ln -s /usr/lib/x86_64-linux-gnu/libtinfo.so.6 /opt/conda/lib/libtinfo.so.6
Load Diff This file was deleted.
You can’t perform that action at this time.
0 commit comments