Upgrade to torch 2.3.1 (#5)

.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.0']  # Must be the most recent version that meets requirements-cuda.txt.
+          pytorch-version: ['2.3.1']  # Should be synced with setup.py.
           cuda-version: ['12.1']
 
     steps:

build.sh

@@ -1,14 +1,9 @@
 #!/bin/bash
+# A simple build script for local testing.
+# NOTE: This script is not used for the actual build process.
 
-eval "$(conda shell.bash hook)"
+PYTORCH_VERSION="2.3.1"
 
-PYTORCH_VERSION="2.3.0"
-
-for PYTHON_VERSION in 38 39 310 311; do
-    source ~/.bashrc;
-    conda activate vllm-flash-py${PYTHON_VERSION};
-    conda env list;
-    pip install packaging ninja;
-    pip install torch==${PYTORCH_VERSION};
-    time python setup.py bdist_wheel --dist-dir=dist;
-done
+pip install packaging ninja;
+pip install torch==${PYTORCH_VERSION};
+time python setup.py bdist_wheel --dist-dir=dist;

setup.py

@@ -263,7 +263,7 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
 
-PYTORCH_VERSION = "2.3.0"
+PYTORCH_VERSION = "2.3.1"
 CUDA_VERSION = "12.1"
 
 setup(