Update requirements

2 years ago · c6fa514b4c
parent e16b10a068
commit c6fa514b4c
2 changed files with 3 additions and 1 deletions
--- a/requirements.txt
+++ b/requirements.txt
@ -9,4 +9,4 @@ transformers>=4.22.0
 torch>=1.10.0
 tqdm>=4.63.0
 cpm_kernels
-deepspeed
+deepspeed>0.6.1
--- a/setup.py
+++ b/setup.py
@ -19,6 +19,8 @@ setup(
        "tokenizers>=0.11.0,<=0.11.4",
        "torch>=1.10.0",
        "tqdm>=4.63.0",
+        "cpm_kernels",
+        "deepspeed>0.6.1",
    ],
    entry_points={}
 )