diff --git a/images/tpu/pytorch/Dockerfile.x86_64 b/images/tpu/pytorch/Dockerfile.x86_64
index 7da256d9ab..bdd00356cb 100644
--- a/images/tpu/pytorch/Dockerfile.x86_64
+++ b/images/tpu/pytorch/Dockerfile.x86_64
@@ -1,3 +1,5 @@
 FROM python:3.10
 
-RUN pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
\ No newline at end of file
+RUN pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
+
+RUN git clone --depth=1 --branch r2.0 https://github.com/pytorch/xla.git
\ No newline at end of file