diff --git a/.github/workflows/slurm_job_scripts/run_e2e_t5x_tests.sub b/.github/workflows/slurm_job_scripts/run_e2e_t5x_tests.sub
index 15755a4cfddd..74ff96d71aca 100644
--- a/.github/workflows/slurm_job_scripts/run_e2e_t5x_tests.sub
+++ b/.github/workflows/slurm_job_scripts/run_e2e_t5x_tests.sub
@@ -53,6 +53,7 @@ rm -rf ${E2E_TESTS_WORKSPACE_DIR}/* \
 && python3.8 -m pip uninstall -y cudf dask-cudf \
 && python3.8 -m pip install ${T5X_DIR} \
 && python3.8 -m pip install ${T5X_DIR}/t5x/contrib/gpu/scripts_gpu/dummy_wikipedia_config \
+&& python3.8 -m pip install -U git+https://github.com/google-research/jestimator \
 && hostname > ${E2E_TESTS_WORKSPACE_DIR}/hostname.txt
 EOF