From d39e002fa5bb758565fb54352ab646b96d922606 Mon Sep 17 00:00:00 2001
From: David Corvoysier <david@huggingface.co>
Date: Thu, 20 Feb 2025 10:13:49 +0000
Subject: [PATCH] feat(neuron): avoid installing CUDA in image

---
 Dockerfile.neuron | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/Dockerfile.neuron b/Dockerfile.neuron
index 5dafc3552..17d256915 100644
--- a/Dockerfile.neuron
+++ b/Dockerfile.neuron
@@ -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 \