add TORCH_VERSION and TODO

RodrigoGonzalez · May 16, 2022 · ba96747 · ba96747
1 parent 2105bc6
commit ba96747
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tpu/Dockerfile b/tpu/Dockerfile
@@ -1,6 +1,7 @@
 ARG BASE_IMAGE_TAG
 ARG LIBTPU_IMAGE_TAG
 ARG TENSORFLOW_VERSION
+ARG TORCH_VERSION
 
 FROM gcr.io/cloud-tpu-v2-images/libtpu:${LIBTPU_IMAGE_TAG} as libtpu
 FROM gcr.io/kaggle-images/python-tpu-tensorflow-whl:python-${BASE_IMAGE_TAG}-${TENSORFLOW_VERSION} AS tensorflow_whl
@@ -17,8 +18,8 @@ RUN pip install /tmp/tensorflow_pkg/tensorflow*.whl && \
 
 # https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm#changing_pytorch_version
 RUN pip uninstall -y torch && \
- pip install torch==1.10 && \
- pip install torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.10-cp37-cp37m-linux_x86_64.whl && \
+ pip install torch==${TORCH_VERSION} && \
+ pip install torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION}-cp37-cp37m-linux_x86_64.whl && \
  /tmp/clean-layer.sh
 
 # https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm

diff --git a/tpu/config.txt b/tpu/config.txt
@@ -1,3 +1,5 @@
+# TODO(b/213335159): Use ci-pretest for BASE_IMAGE_TAG once stable.
 BASE_IMAGE_TAG=v108
 LIBTPU_IMAGE_TAG=libtpu_1.1.0_RC00
-TENSORFLOW_VERSION=2.8.0
+TENSORFLOW_VERSION=2.8.0
+TORCH_VERSION=1.11.0