fixing an issue in pytorch docker images with cuda failure

KalanaRatnayake · Aug 8, 2024 · 12cafd8 · 12cafd8
1 parent 4a7c017
commit 12cafd8
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/pytorch-images/r3630.Dockerfile b/pytorch-images/r3630.Dockerfile
@@ -15,11 +15,21 @@ WORKDIR /
 RUN apt-get update -y
 RUN apt-get install -y --no-install-recommends python3-pip
 
+#####################################################################################
+##                           Install PyTorch 2.4
+#####################################################################################
+
+RUN wget https://developer.download.nvidia.cn/compute/redist/jp/v60/pytorch/torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl
+
+RUN python3 -m pip install --no-cache-dir torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl
+
+RUN rm torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl
+
 #####################################################################################
 ##                           Install PyTorch
 #####################################################################################
 
-RUN python3 -m pip install --no-cache-dir torch torchvision torchaudio
+RUN python3 -m pip install --no-cache-dir torchvision torchaudio
 
 #####################################################################################
 ##