diff --git a/requirements_mtj.txt b/requirements_mtj.txt
index 759e3b1c..076d9429 100644
--- a/requirements_mtj.txt
+++ b/requirements_mtj.txt
@@ -6,6 +6,7 @@ dm-haiku == 0.0.5
 jax == 0.2.21
 jaxlib >= 0.1.69, <= 0.3.7
 transformers == 4.25.1
+chex == 0.1.5
 huggingface_hub==0.12.1
 progressbar2
 git+https://github.com/VE-FORBRYDERNE/mesh-transformer-jax@ck