diff --git a/.github/workflows/docker-rocm.yaml b/.github/workflows/docker-rocm.yaml
index 17a3f03f..690ec8bf 100644
--- 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