diff --git a/.github/workflows/python_tests.yml b/.github/workflows/python_tests.yml
index 092753719..fb5b85f2b 100644
--- a/.github/workflows/python_tests.yml
+++ b/.github/workflows/python_tests.yml
@@ -148,9 +148,15 @@ jobs:
         TAG=macosx_10_15_x86_64.macosx_11_0_x86_64.macosx_12_0_x86_64
         python tests/ci_build/rename_whl.py python-package/dist/*.whl ${{ github.sha }} ${TAG}
 
+    - name: Extract branch name
+      shell: bash
+      run: echo "##[set-output name=branch;]$(echo ${GITHUB_REF#refs/heads/})"
+      id: extract_branch
+      if: github.ref == 'refs/heads/master' || contains(github.ref, 'refs/heads/release_')
+
     - name: Upload Python wheel
       shell: bash -l {0}
-      if: (github.ref == 'refs/heads/master' || contains(github.ref, 'refs/heads/release_'))
+      if: github.ref == 'refs/heads/master' || contains(github.ref, 'refs/heads/release_')
       run: |
         python -m awscli s3 cp python-package/dist/*.whl s3://xgboost-nightly-builds/${{ steps.extract_branch.outputs.branch }}/ --acl public-read
       env: