diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
index a934401a20..d95a440f55 100644
--- a/.github/workflows/test_cuda.yml
+++ b/.github/workflows/test_cuda.yml
@@ -19,7 +19,7 @@ jobs:
     runs-on: nvidia
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
     container:
-      image: nvidia/cuda:12.6.2-cudnn-devel-ubuntu22.04
+      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:
@@ -36,12 +36,6 @@ jobs:
       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 \
-         && 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
     - 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: |
diff --git a/pyproject.toml b/pyproject.toml
index 5ad5101c01..7be0554070 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -157,7 +157,7 @@ pin_tensorflow_cpu = [
   "tensorflow; platform_machine=='x86_64' and platform_system == 'Darwin'",
 ]
 pin_tensorflow_gpu = [
-  "tensorflow~=2.18.0",
+  "tensorflow==2.19.1",
 ]
 pin_pytorch_cpu = [
   # https://github.com/pytorch/pytorch/issues/114602
@@ -169,7 +169,7 @@ pin_pytorch_gpu = [
   "torch>=2.7,<2.10",
 ]
 pin_jax = [
-  "jax==0.5.0;python_version>='3.10'",
+  "jax==0.6.2;python_version>='3.10'",
 ]
 
 [tool.setuptools_scm]