diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 37c284888..8e677398b 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -49,7 +49,7 @@ jobs:
       matrix:
           os: ['ubuntu-20.04']
           python-version: ['3.8', '3.9', '3.10', '3.11']
-          pytorch-version: ['2.3.1']  # Should be synced with setup.py.
+          pytorch-version: ['2.4.0']  # Should be synced with setup.py.
           cuda-version: ['12.1']
 
     steps:
diff --git a/build.sh b/build.sh
index 78fcd0a62..45a8eb281 100755
--- a/build.sh
+++ b/build.sh
@@ -2,7 +2,7 @@
 # A simple build script for local testing.
 # NOTE: This script is not used for the actual build process.
 
-PYTORCH_VERSION="2.3.1"
+PYTORCH_VERSION="2.4.0"
 
 pip install packaging ninja;
 pip install torch==${PYTORCH_VERSION};
diff --git a/setup.py b/setup.py
index 4c6142bca..b3b7fef96 100644
--- a/setup.py
+++ b/setup.py
@@ -263,7 +263,7 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
 
-PYTORCH_VERSION = "2.3.1"
+PYTORCH_VERSION = "2.4.0"
 CUDA_VERSION = "12.1"
 
 setup(