add TORCH_VERSION and TODO

Author: djherbis

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

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