Merge pull request #1289 from Kaggle/add-tpu-pytorch-deps

Add TPU deps from PyTorch example
Kaggle · Aug 15, 2023 · 3bb531e · 3bb531e
2 parents bbf8424 + 79b7bd4
commit 3bb531e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tpu/Dockerfile b/tpu/Dockerfile
@@ -52,7 +52,7 @@ RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-ar
     jax[tpu]==${JAX_VERSION} -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax elegy git+https://github.com/deepmind/dm-haiku jraph distrax \
     numpy==1.23.5 \
     papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \
-    pandas matplotlib opencv-python-headless librosa scikit-learn
+    pandas matplotlib opencv-python-headless librosa scikit-learn accelerate diffusers transformers
 
 # Tensorflow libtpu:
 RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so