diff --git a/requirements.txt b/requirements.txt
index 6af50f127..bffaf892f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,11 +6,11 @@ pymanopt
 cvxopt
 scikit-learn
 torch
-jax
-jaxlib
+jax<=0.4.24
+jaxlib<=0.4.24
 tensorflow
 pytest
 torch_geometric
 cvxpy
 geomloss
-pykeops
\ No newline at end of file
+pykeops