diff --git a/requirements_dev.txt b/requirements_dev.txt
index 047f20a4cc4..dcbfe8f6dad 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -27,7 +27,7 @@ docker
 stable_baselines3
 tensorboard
 gym
-jax[cpu]; sys_platform == 'darwin' or sys_platform == 'linux'
+jax[cpu]<0.4; sys_platform == 'darwin' or sys_platform == 'linux'
 fastcore; python_version > '3.6'
 fastcore==1.3.29; python_version == '3.6'
 pyarrow