Fix PyTorch installation in ROCm training image

opendatahub-io · Nov 18, 2024 · b2a7320 · b2a7320
1 parent 67102bf
commit b2a7320
Show file tree

Hide file tree

Showing 2 changed files with 657 additions and 683 deletions.
diff --git a/images/runtime/training/rocm/Pipfile b/images/runtime/training/rocm/Pipfile
@@ -4,7 +4,7 @@ verify_ssl = true
 name = "pypi"
 
 [[source]]
-url = "https://download.pytorch.org/whl/rocm6.1"
+url = "https://download.pytorch.org/whl/"
 verify_ssl = true
 name = "pytorch"
 
@@ -14,7 +14,8 @@ datasets = ">=2.15.0"
 transformers = ">=4.41.2"
 numpy = "<2.0.0,>=1.23.5"
 accelerate = ">=0.34.2"
-torch = {version = "==2.4.1", index = "pytorch"}
+torch = {version = "==2.4.1+rocm6.1", index = "pytorch"}
+pytorch-triton-rocm = {version = "*", index = "pytorch"}
 sentencepiece = "<0.3,>=0.1.99"
 tokenizers = "<1.0,>=0.13.3"
 tqdm = "<5.0,>=4.66.2"