fix: checkpoint saving with deepspeed (#1321)

2024-02-27 15:46:44 +09:00
parent 0f6af36d50
commit 5be8b555a0
2 changed files with 2 additions and 2 deletions
--- a/setup.py
+++ b/setup.py
@@ -74,7 +74,7 @@ setup(
            "fused-dense-lib  @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",
        ],
        "deepspeed": [
-            "deepspeed>=0.13.1",
+            "deepspeed==0.13.1",
            "deepspeed-kernels",
        ],
        "mamba-ssm": [