Update torch to 2.4.0 (#8)

Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

.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:

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};

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(