Fix rocm ci

pytorch-labs · Nov 22, 2024 · f1f9b3b · f1f9b3b
1 parent e0c459e
commit f1f9b3b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/workflows/docker-rocm.yaml b/.github/workflows/docker-rocm.yaml
@@ -48,7 +48,7 @@ jobs:
         if: github.event_name != 'pull_request'
         run: |
           # Extract pytorch version from the docker
-          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
+          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:rocm-latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
           export DOCKER_TAG=$(awk '{match($0, /dev[0-9]+/, arr); print arr[0]}' <<< "${PYTORCH_VERSION}")
           docker push ghcr.io/pytorch-labs/tritonbench:rocm-${DOCKER_TAG}
           docker push ghcr.io/pytorch-labs/tritonbench:rocm-latest