fix: restore sm_60 (P100) PyTorch kernels in GPU image (best-effort, unverified against real fleet)

[shivansh193]

What

Fixes #1546 (Tesla P100 / sm_60 + recent PyTorch build incompatibility). For GPU images only, reinstalls the same torch/torchvision/torchaudio version already selected by the Colab base image, but from the cu126 wheel index instead of whichever CUDA index the base image used (currently cu128).

Root cause (as I found it)

• The GPU image is FROM us-docker.pkg.dev/colab-images/public/runtime:.... Dockerfile.tmpl doesn't install PyTorch directly — it pip freezes the base image's existing torch/torchvision/etc. versions into /colab_requirements.txt and reinstalls those exact versions (see lines 9-17). So the actual PyTorch wheel/CUDA-index choice is inherited from the upstream Colab image, not controlled by anything in this repo today.
• Per the issue thread and my own check of pytorch/pytorch's current .ci/manywheel/build_env_setup.py (TORCH_CUDA_ARCH_LIST_TABLE), PyTorch's officially published wheels have per-CUDA-version compute-capability sets, e.g.:
  • cu126 → {50, 60, 70, 75, 80, 86, 90} (x86_64) — includes sm_60
  • newer CUDA lines (matching the sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120 reported in the issue and in torch.cuda's own startup warning) drop sm_50/sm_60 in exchange for sm_100/sm_120 (Blackwell) support.
• This matches multiple comments on Pytorch CUDA P100 GPU Incompatibility #1546: reinstalling the same torch version from --index-url https://download.pytorch.org/whl/cu126 was independently found by users to fix P100, while leaving T4 (sm_75) unaffected.
• I also noticed tests/common.py already has a long-standing p100_exempt test decorator (b/342143152 P100s are slowly being unsupported in new release of popular ml tools such as RAPIDS), and tests/test_pytorch.py's p100_exempt-marked tests (test_gpu_computation, a .sum() reduction; test_linalg) line up exactly with the failure mode reported in Pytorch CUDA P100 GPU Incompatibility #1546 and its comments. So this appears to be a recurring, known category of breakage rather than a one-off regression, which is part of why I think a build-level fix is worth proposing rather than only a per-user workaround.

The change

In Dockerfile.tmpl, right after the existing "Install Kaggle packages" step (which installs whatever torch version the base image already has), add a GPU-only step that force-reinstalls the same torch/torchvision/torchaudio version from the cu126 index:

{{ if eq .Accelerator "gpu" }}
RUN TORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])") && \
    uv pip install --system --no-cache --force-reinstall \
        "torch==${TORCH_VERSION}" torchvision torchaudio \
        --index-url https://download.pytorch.org/whl/cu126 \
        --extra-index-url https://pypi.org/simple
{{ end }}

I deliberately did not hardcode a torch version — it reads back whatever version is already installed (from the base image) and re-pulls that same version's cu126 build, so this doesn't fight the existing "freeze the base image's version" mechanism and shouldn't need to be bumped every time the Colab base image updates. It's scoped to the GPU template branch only; I rendered Dockerfile.tmpl locally with renderizer --ACCELERATOR=gpu and --ACCELERATOR=none to confirm the new block only appears in the GPU output.

What I could not verify

I don't have access to Kaggle's build pipeline, its Jenkins GPU executors, or an actual P100/T4 to test against, so I could not:

• Actually build this Dockerfile against the real base image and confirm the install succeeds without dependency conflicts (e.g. other packages in kaggle_requirements.txt or the base image that may have been compiled/linked against the cu128 build's ABI — I don't have visibility into whether anything like that exists).
• Confirm on real hardware that this restores P100 functionality end-to-end (only that PyTorch's own published cu126 wheel arch table includes sm_60, and that this specific cu128→cu126 swap-at-same-version was independently reported to work by commenters on Pytorch CUDA P100 GPU Incompatibility #1546).
• Confirm there isn't a newer/better-supported PyTorch release line that covers both sm_60 and sm_100/120 simultaneously — as far as I could tell from PyTorch's current build matrix, that combination isn't offered in any single official wheel today, so this is a "pick one" tradeoff, not a strict win.

So please treat this as a best-effort, plausible starting point for someone with access to your actual CI/build pipeline to verify and adjust, not a guaranteed-correct fix.

Other context I noticed while investigating

While looking through recent branches/PRs I came across #1560 ("chore: drop P100 from the CI pipeline"), which is tearing down the P100 Jenkins CI/build agent. I don't know if that reflects a broader decision to deprecate P100 as a user-facing accelerator too, in which case this PR may be moot — flagging it in case it's useful context for triage, and apologies in advance if this duplicates effort already in flight internally.

Scope

This only touches Dockerfile.tmpl's GPU branch. I did not touch tests, CI, or anything else, per the intent of keeping this change minimal and targeted at the specific issue.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[calderjo]

Hey Kaggle staff here, i saw that you mention a PR i am working on

note that google cloud is deprecating p100s which Kaggle relies on:
https://www.kaggle.com/discussions/product-announcements/735239

Kaggle team is working on moving users over to T4 for new notebooks and sessions,
and we have no intention of maintaining p100 support in future docker image releases.