huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-19 22:02:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(neuron): avoid installing CUDA in image

Browse Source
This commit is contained in:
David Corvoysier 2025-02-20 10:13:49 +00:00 committed by Nicolas Patry
parent 7d6ff64c13
commit d39e002fa5
No known key found for this signature in database
GPG Key ID: 4242CEF24CB6DBF9
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
Dockerfile.neuron
View File

@ -120,6 +120,12 @@ RUN apt-get update -y \
ENV PATH="/opt/bin/:/opt/aws/neuron/bin:${PATH}"
# Install manually torch CPU version to avoid pulling CUDA
RUN pip3 install \
torch==2.1.2 \
torchvision==0.16.2 \
--index-url https://download.pytorch.org/whl/cpu
RUN pip3 install \
neuronx-cc==2.15.143.0 \
torch-neuronx==2.1.2.2.3.2 \