let try this out

Author: calderjo

tpu/Dockerfile

@@ -25,36 +25,13 @@ ADD patches/kaggle_session.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packa
 ADD patches/kaggle_web_client.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/kaggle_web_client.py
 ADD patches/kaggle_datasets.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/kaggle_datasets.py
 
-# Disable GCP integrations for now
-# ADD patches/kaggle_gcp.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/kaggle_gcp.py
-
-# Disable logging to file (why do we need this?)
-# ADD patches/log.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/log.py
-
-# sitecustomize adds significant latency to ipython kernel startup and should only be added if needed
-# ADD patches/sitecustomize.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/sitecustomize.py
-
 # Prereqs
 # This is needed for cv2 (opencv-python):
 # https://stackoverflow.com/questions/55313610/importerror-libgl-so-1-cannot-open-shared-object-file-no-such-file-or-directo
 RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
 
 # Install all the packages together for maximum compatibility.
-
-# Install Tensorflow.
-
-# Install Pytorch & related packages
-# https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm#changing_pytorch_version
-# The URL doesn't include patch version. i.e. must use 1.11 instead of 1.11.0
-# We need to keep the numpy version the same as the installed tf one but compatible with other installs.
-
-# Install JAX & related packages
-# https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm
-
-# Packages needed by the Notebook editor
-
 # Additional useful packages should be added in the requirements.txt
-
 # Bring in the requirements.txt and replace variables in it:
 RUN apt-get install -y gettext
 ADD tpu/requirements.txt /kaggle_requirements.txt

tpu/config.txt

@@ -1,12 +1,12 @@
-BASE_IMAGE=python:3.10
-PYTHON_WHEEL_VERSION=cp310
-PYTHON_VERSION_PATH=python3.10
-TENSORFLOW_VERSION=2.18.0
+BASE_IMAGE=python:3.11
+PYTHON_WHEEL_VERSION=cp311
+PYTHON_VERSION_PATH=python3.11
+TENSORFLOW_VERSION=2.19.0
 # gsutil ls gs://pytorch-xla-releases/wheels/tpuvm/* | grep libtpu | grep torch_xla | grep -v -E ".*rc[0-9].*" | sed 's/.*torch_xla-\(.*\)+libtpu.*/\1/' | sort -rV
 # Supports nightly
-TORCH_VERSION=2.5.0
+TORCH_VERSION=2.6.0
 # https://github.com/pytorch/audio supports nightly
-TORCHAUDIO_VERSION=2.5.0
+TORCHAUDIO_VERSION=2.6.0
 # https://github.com/pytorch/vision supports nightly
-TORCHVISION_VERSION=0.20.0
+TORCHVISION_VERSION=0.21.0
 TORCH_LINUX_WHEEL_VERSION=manylinux_2_28_x86_64