Add tensor2tensor to Dockerfile since seqio doesn't explicitly depend…

… on it. PiperOrigin-RevId: 454974059
google · Jun 14, 2022 · 7e864af · 7e864af
1 parent f756968
commit 7e864af
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/paxml/pip_package/Dockerfile b/paxml/pip_package/Dockerfile
@@ -78,7 +78,8 @@ RUN cd paxml && bazel build paxml/tasks/lm/params:main
 RUN pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 RUN pip3 install fiddle-config tensorstore
 RUN pip3 uninstall -y seqio seqio-nightly
-RUN pip3 install -U seqio-nightly
+# Need tensor2tensor explicitly because seqio implicitly dep on tensor2tensor.
+RUN pip3 install -U seqio-nightly tensor2tensor
 
 WORKDIR "/paxml"