diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
index 4be035ba21..7b46563d11 100644
--- a/.github/workflows/test_cuda.yml
+++ b/.github/workflows/test_cuda.yml
@@ -18,30 +18,34 @@ jobs:
     name: Test Python and C++ on CUDA
     runs-on: nvidia
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
-    container:
-      image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
-      options: --gpus all
+    # container:
+    #   image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
+    #   options: --gpus all
     if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
-      - name: Make sudo and git work
-        run: apt-get update && apt-get install -y sudo git
       - uses: actions/checkout@v6
       - uses: actions/setup-python@v6
         with:
           python-version: "3.11"
           # cache: 'pip'
       - name: Install wget and unzip
-        run: apt-get update && apt-get install -y wget unzip
+        run: sudo apt-get update && sudo apt-get install -y wget unzip
       - uses: lukka/get-cmake@latest
         with:
           useLocalCache: true
           useCloudCache: false
       - run: |
-          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
-          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
+          UBUNTU_VERSION=$(lsb_release -rs | tr -d '.')
+          OS_NAME="ubuntu${UBUNTU_VERSION}"
+          echo "Current OS: ${OS_NAME}"
+
+          wget https://developer.download.nvidia.com/compute/cuda/repos/${OS_NAME}/x86_64/cuda-keyring_1.1-1_all.deb \
+          && sudo dpkg -i cuda-keyring_1.1-1_all.deb \
           && sudo apt-get update \
-          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
-        if: false # skip as we use nvidia image
+          && sudo apt-get -y install cuda-toolkit-12-9 cudnn9-cuda-12
+          echo "CUDA_PATH=/usr/local/cuda-12.9" >> $GITHUB_ENV
+          echo "/usr/local/cuda-12.9/bin" >> $GITHUB_PATH
+        # 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_gpu
       - run: |