diff --git a/requirements/requirements-core.txt b/requirements/requirements-core.txt
index 7b4c342..0af3598 100644
--- a/requirements/requirements-core.txt
+++ b/requirements/requirements-core.txt
@@ -12,6 +12,7 @@ tqdm
 jax
 jaxlib
 flax
+#chex==0.1.6
 distrax
 cloudpickle
 getdist
\ No newline at end of file