Use index-url rather then extra-index-url

pytorch · Feb 26, 2024 · a46615f · a46615f
1 parent bf7fd30
commit a46615f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/.github/workflows/test_linux_job.yml b/.github/workflows/test_linux_job.yml
@@ -48,7 +48,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch
         python3 -c 'import torch'
   test-gpu:
@@ -57,24 +57,24 @@ jobs:
       matrix:
         runner_type: ["linux.4xlarge.nvidia.gpu", "linux.g5.4xlarge.nvidia.gpu"]
     with:
-      job-name: "linux-py3.8-cu116"
+      job-name: "linux-py3.8-cu121"
       runner: ${{ matrix.runner_type }}
       test-infra-repository: ${{ github.repository }}
       test-infra-ref: ${{ github.ref }}
       submodules: ${{ 'true' }}
       gpu-arch-type: cuda
-      gpu-arch-version: "11.6"
+      gpu-arch-version: "12.1"
       timeout: 60
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cu116 --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cu121 --pre torch
         # Can import pytorch, cuda is available
         python3 -c 'import torch;cuda_avail = torch.cuda.is_available();print("CUDA available: " + str(cuda_avail));assert(cuda_avail)'
         python3 -c 'import torch;t = torch.ones([2,2], device="cuda:0");print(t);print("tensor device:" + str(t.device))'
         nvidia-smi
-        nvcc --version | grep "cuda_11.6"
-        [[ "${CUDA_HOME}" == "/usr/local/cuda-11.6" ]] || exit 1
+        nvcc --version | grep "cuda_12.1"
+        [[ "${CUDA_HOME}" == "/usr/local/cuda-12.1" ]] || exit 1
   test-docker-image:
     uses: ./.github/workflows/linux_job.yml
     with:

diff --git a/tools/pkg-helpers/pytorch_pkg_helpers/wheel.py b/tools/pkg-helpers/pytorch_pkg_helpers/wheel.py
@@ -24,7 +24,7 @@ def get_pytorch_pip_install_command(
     if channel == "nightly":
         pip_install += " --pre"
     extra_index = f"https://download.pytorch.org/whl/{channel}/{gpu_arch_version}"
-    return [f"export PIP_INSTALL_TORCH='{pip_install} --extra-index-url {extra_index}'"]
+    return [f"export PIP_INSTALL_TORCH='{pip_install} --index-url {extra_index}'"]
 
 
 def get_pytorch_s3_bucket_path(