CUDA, drivers and container images

Setup#

pip install torch

Checking your own stack, for real#

import subprocess
import torch

driver_out = subprocess.run(
    ["nvidia-smi", "--query-gpu=driver_version,name", "--format=csv,noheader"],
    capture_output=True, text=True, timeout=5,
).stdout.strip()
driver_version, gpu_name = [s.strip() for s in driver_out.split(",")]

print(f"GPU:                        {gpu_name}")
print(f"NVIDIA driver version:      {driver_version}")
print(f"PyTorch was built for CUDA: {torch.version.cuda}")
print(f"PyTorch version:            {torch.__version__}")
print(f"CUDA available to PyTorch:  {torch.cuda.is_available()}")

if torch.cuda.is_available():
    x = torch.randn(3, 3, device="cuda")
    y = x @ x
    print(f"a real matrix multiply on the GPU worked: result shape {tuple(y.shape)}")

GPU:                        NVIDIA RTX A6000
NVIDIA driver version:      556.39
PyTorch was built for CUDA: 12.1
PyTorch version:            2.5.1+cu121
CUDA available to PyTorch:  True
a real matrix multiply on the GPU worked: result shape (3, 3)

Real output from this machine's actual GPU stack. Yours will show different version numbers, and that is completely fine. Check not for any specific number, but that CUDA available to PyTorch reads True, and the matrix multiply line actually appears — confirming the GPU genuinely did the work.

Line-by-line walkthrough#

nvidia-smi --query-gpu=driver_version. Asks the driver directly for its own version. That is the ground truth for what the machine actually has installed, independent of anything Python or PyTorch reports.

torch.version.cuda. This is the CUDA version PyTorch was built against when it was compiled. Not the CUDA version installed on your machine. A common point of confusion: this number can differ from what nvcc --version reports on the machine itself. That is expected, not a bug.

torch.cuda.is_available(). The single most important line for catching a mismatch. Say this is False on a machine that visibly has a working GPU, visible to nvidia-smi. The driver and the installed PyTorch build do not agree. The reason is almost always a version mismatch somewhere in the stack.

A matching Docker base image, for reference#

# The CUDA version in this base image tag must be supported by the HOST
# machine's driver -- see the compatibility table below.
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

RUN pip install torch --index-url https://download.pytorch.org/whl/cu121

There is no output block for this. Building and running a GPU-enabled container needs the NVIDIA Container Toolkit installed on the host, which cannot be demonstrated inside a single script. The line to focus on is the version match between the base image's CUDA tag (12.1.0) and the PyTorch build's expected CUDA version (cu121). Both must also be supported by the host machine's actual driver.

Common mistakes#

Assuming a container fixes every compatibility problem. A container bundles its own CUDA toolkit, but it still needs the host machine's driver to support that CUDA version. Containerising code does not remove the driver requirement. It only removes the CUDA toolkit and framework install from the list of things that can mismatch.

Checking torch.__version__ and stopping there. The framework version alone says nothing about whether it can actually reach the GPU. torch.cuda.is_available() is the check that actually matters.

Silently falling back to CPU without noticing. Code that never explicitly checks torch.cuda.is_available() before running can execute correctly on the CPU for a long time, before anyone notices it never touched the GPU at all. Correct answers, at a small fraction of the expected speed.

Mixing driver versions across a fleet of machines. A model tested and confirmed working on one machine can silently fail to use the GPU on another machine in the same fleet. That happens if that machine's driver was never updated to match. Check this per machine, not once for the whole fleet.

Try it yourself#

Run nvcc --version in a terminal, if it is installed, and compare its reported CUDA version against torch.version.cuda. They are allowed to differ — understanding why they are allowed to differ is the point of the exercise.

What to learn next#

• Docker for ML — the packaging layer this lesson's compatibility rules apply inside.
• Serverless GPU platforms — where this whole compatibility stack is managed for you, for a price.
• Moving tensors between CPU and GPU — the everyday code most affected when this compatibility silently breaks.