Install triton package while running triton tests

pytorch · May 10, 2024 · c616e64 · c616e64
1 parent 256d819
commit c616e64
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/.github/workflows/_triton.yml b/.github/workflows/_triton.yml
@@ -111,6 +111,7 @@ jobs:
           pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
           pip install -U --pre jax-cuda12-pjrt jax-cuda12-plugin -f https://storage.googleapis.com/jax-releases/jax_cuda_plugin_nightly_releases.html
           pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+          pip install --no-deps triton==2.3.0
       - name: Run Tests
         env:
           PJRT_DEVICE: CUDA