diff --git a/setup.py b/setup.py index df483c1ec9c..48250a5634f 100644 --- a/setup.py +++ b/setup.py @@ -113,8 +113,8 @@ _deps = [ "importlib_metadata", "ipadic>=1.0.0,<2.0", "isort>=5.5.4", - "jax>=0.2.8,!=0.3.2", - "jaxlib>=0.1.65", + "jax>=0.2.8,!=0.3.2,<=0.3.6", + "jaxlib>=0.1.65,<=0.3.6", "jieba", "nltk", "numpy>=1.17", diff --git a/src/transformers/dependency_versions_table.py b/src/transformers/dependency_versions_table.py index 334103c20a5..2468fe9b6e1 100644 --- a/src/transformers/dependency_versions_table.py +++ b/src/transformers/dependency_versions_table.py @@ -23,8 +23,8 @@ deps = { "importlib_metadata": "importlib_metadata", "ipadic": "ipadic>=1.0.0,<2.0", "isort": "isort>=5.5.4", - "jax": "jax>=0.2.8,!=0.3.2", - "jaxlib": "jaxlib>=0.1.65", + "jax": "jax>=0.2.8,!=0.3.2,<=0.3.6", + "jaxlib": "jaxlib>=0.1.65,<=0.3.6", "jieba": "jieba", "nltk": "nltk", "numpy": "numpy>=1.17",