diff --git a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
index 9272f6e627a13..4e55ce29f46ff 100644
--- a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
+++ b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
@@ -125,7 +125,8 @@ RUN pip install \
     sentencepiece \
     wget \
     dill==0.3.4 \
-    pytorch_lightning==1.6.0 \
+    pytorch_lightning==2.3.3 \
+    tensorboard \
     pytest-xdist \
     pytest-rerunfailures \
     ml_dtypes==0.3.0 \