ci: pin PT to 2.3.1 when using CUDA (#4009)

PT 2.4.0 requires cudnn 9, incompatible with the latest TF with cudnn 8.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Updated the build process to optimize resource usage by removing
unnecessary files.
- Specified a fixed version for PyTorch to ensure consistent
functionality across environments.
	
- **Documentation**
- Added configuration settings to manage library dependencies and ensure
compatibility.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

.github/workflows/build_wheel.yml

@@ -142,6 +142,8 @@ jobs:
           - variant: "_cu11"
             cuda_version: "11"
     steps:
+      - name: Delete huge unnecessary tools folder
+        run: rm -rf /opt/hostedtoolcache
       - uses: actions/checkout@v4
       - uses: actions/download-artifact@v4
         with:

.github/workflows/test_cuda.yml

@@ -47,7 +47,7 @@ jobs:
          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
     - run: python -m pip install -U uv
-    - run: source/install/uv_with_retry.sh pip install --system "tensorflow>=2.15.0rc0" "torch>=2.2.0"
+    - run: source/install/uv_with_retry.sh pip install --system "tensorflow>=2.15.0rc0" "torch==2.3.1.*"
     - run: |
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')

pyproject.toml

@@ -266,6 +266,8 @@ PATH = "/usr/lib64/mpich/bin:$PATH"
 UV_EXTRA_INDEX_URL = "https://download.pytorch.org/whl/cpu"
 # trick to find the correction version of mpich
 CMAKE_PREFIX_PATH="/opt/python/cp311-cp311/"
+# PT 2.4.0 requires cudnn 9, incompatible with TF with cudnn 8
+PYTORCH_VERSION = "2.3.1"
 
 [tool.cibuildwheel.windows]
 test-extras = ["cpu", "torch"]