How to Build a CUDA Docker Image in GitHub Actions

Base the image on an nvidia/cuda tag that matches your torch build, install the runtime, and build it in CI so every run uses the same GPU-ready container.

Start from an nvidia/cuda runtime base whose CUDA version matches your torch wheel, install Python and dependencies, then build with docker/build-push-action. The base ships the CUDA runtime libraries torch needs at load time.

Dockerfile

FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04
RUN apt-get update && apt-get install -y python3 python3-pip && rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY requirements.txt .
RUN pip3 install --no-cache-dir torch==2.3.1 --index-url https://download.pytorch.org/whl/cu121 \
 && pip3 install --no-cache-dir -r requirements.txt
COPY . .
CMD ["python3", "train.py"]

Workflow

jobs:
  build:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: docker/build-push-action@v6
        with:
          context: .
          tags: my-org/trainer:${{ github.sha }}

Gotchas

• The CUDA base tag must match the torch cu version, or torch fails to load libcudart at runtime.
• Use a runtime base for inference; the devel base is only needed to compile CUDA code.

Related guides