From b0a1d667b04dcce8e54201cb429a57155bf9670c Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Sun, 17 Dec 2023 01:46:54 -0800
Subject: [PATCH] Pin PyTorch & xformers versions (#2155)

---
 .github/workflows/publish.yml | 2 +-
 pyproject.toml                | 2 +-
 requirements-build.txt        | 2 +-
 requirements.txt              | 4 ++--
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index a7d60150..5211dc18 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.1.1']
+          pytorch-version: ['2.1.2']  # Must be the most recent version that meets requirements.txt.
           cuda-version: ['11.8', '12.1']
 
     steps:
diff --git a/pyproject.toml b/pyproject.toml
index d18a8774..b197256f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ requires = [
     "ninja",
     "packaging",
     "setuptools >= 49.4.0",
-    "torch >= 2.1.1",
+    "torch == 2.1.2",
     "wheel",
 ]
 build-backend = "setuptools.build_meta"
diff --git a/requirements-build.txt b/requirements-build.txt
index 38514d27..7e7e48a1 100644
--- a/requirements-build.txt
+++ b/requirements-build.txt
@@ -2,5 +2,5 @@
 ninja
 packaging
 setuptools>=49.4.0
-torch>=2.1.0
+torch==2.1.2
 wheel
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 21473420..f509bdf4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,9 +5,9 @@ pandas  # Required for Ray data.
 pyarrow  # Required for Ray data.
 sentencepiece  # Required for LLaMA tokenizer.
 numpy
-torch >= 2.1.1
+torch == 2.1.2
 transformers >= 4.36.0  # Required for Mixtral.
-xformers >= 0.0.23  # Required for CUDA 12.1.
+xformers == 0.0.23  # Required for CUDA 12.1.
 fastapi
 uvicorn[standard]
 pydantic == 1.10.13  # Required for OpenAI server.