set torch_backend to cu128

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

Author: njzjz

.github/workflows/test_cuda.yml

@@ -43,7 +43,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 --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax "jax[cuda12]"
+      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax "jax[cuda12]" --torch-backend cu128
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')