Merge pull request #71 from Yhlong00/main

fix: incorrect cuda
runpod · Sep 24, 2024 · 521a178 · 521a178
2 parents 83549ee + 3eefd14
commit 521a178
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/official-templates/pytorch/README.md b/official-templates/pytorch/README.md
@@ -5,8 +5,9 @@
 - To specify the platform, use `docker buildx bake <target> --set <target>.platform=linux/amd64`.
 
 Example:
+
 ```bash
-docker buildx bake 240-py311-cuda1240-devel-ubuntu2204 --set 240-py311-cuda1240-devel-ubuntu2204.platform=linux/amd64
+docker buildx bake 240-py311-cuda1241-devel-ubuntu2204 --set 240-py311-cuda1241-devel-ubuntu2204.platform=linux/amd64
 ```
 
 ## Exposed Ports

diff --git a/official-templates/pytorch/docker-bake.hcl b/official-templates/pytorch/docker-bake.hcl
@@ -184,7 +184,7 @@ target "240-py311-cuda1241-devel-ubuntu2204" {
     args = {
         BASE_IMAGE = "nvidia/cuda:12.4.1-devel-ubuntu22.04"
         PYTHON_VERSION = "3.11"
-        TORCH = "torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0"
+        TORCH = "torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124"
     }
 }