fix triton test execution

NVIDIA · Nov 26, 2024 · 3efecc2 · 3efecc2
1 parent 81886ec
commit 3efecc2
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -458,8 +458,9 @@ jobs:
         docker run -i --shm-size=1g --gpus all --volume $PWD:/output \
         ${{ needs.build-triton.outputs.DOCKER_TAG_FINAL }} \
         bash <<"EOF" |& tee test-triton.log
-          # autotuner tests from jax-triton now hit a triton code path that uses utilities from pytorch...
-          pip install --no-deps torch --index-url https://download.pytorch.org/whl/cpu
+          # autotuner tests from jax-triton now hit a triton code path that uses utilities from pytorch; this relies on
+          # actually having a CUDA backend for pytoch
+          pip install --no-deps torch
           python /opt/jax-triton/tests/triton_call_test.py --xml_output_file /output/triton_test.xml
         EOF
       STATISTICS_SCRIPT: |