Lab: Your Python Environment and a First Trained Model

Validated on: written from the documentation cited above; not yet validated on hardware on any track. The versions each track was run with will be recorded here when the validation pass is done.

Objective

Before executing, read the lab execution and evidence guide. Use this lesson's explicit working directories and track setup; keep each server in its own terminal. Record hardware validation as pass, fail or not run, with the evidence requested below.

By the end of this lab you will have, on your own machine:

• a directory ~/llm-course holding a Python virtual environment made by uv, with the accelerator build of PyTorch for your track (CUDA on S and N, ROCm on X, MPS on M, and MLX as well on M), torchvision, JupyterLab and matplotlib, and a one-line proof of which device PyTorch will use;
• the two-layer network from the previous lessons trained on MNIST, the classic set of 70,000 handwritten digits, with the loss printed per epoch, plotted and read, plus one run that overfits on purpose and two that fail on purpose, so that you have seen with your own numbers what the generalisation lesson drew as a picture;
• a checkpoint file whose bytes you can account for to the last tensor, reloaded into a fresh network and evaluated without a single training step, and a JupyterLab notebook, part-01.ipynb, in which you looked at the curves and the weights;
• labbook.md, the lab notebook you keep for the rest of the course, with its Machine and Environment sections filled in and a JSON record of every run.

Nothing here needs a big machine: the model has 203,530 parameters and five epochs take seconds on an accelerator, minutes on a CPU. What needs care is the environment, because every later lab assumes it, and a wheel built for the wrong accelerator is the usual reason a Level 3 training run turns out to be running on the CPU.

⌂Single-machine path

This lab is a single-machine lab on every track. Owners of a two-node Spark cluster or several Macs run it on one machine; the second one is not used until Part 18.

The stack you are about to build

Every framework in this course sits on the same five layers, and each can fail on its own; the lab builds them bottom-up and checks each before moving to the next.

What sits under python train-mnist.py

built and checked in this order

1. Your scripts and notebookstrain-mnist.py, plot-curves.py, reload-mnist.py, part-01.ipynb, and labbook.md recording what they produced.Tasks 4 to 10
2. The virtual environment~/llm-course/.venv: a private Python 3.12 with torch, torchvision, jupyterlab and matplotlib in its site-packages. Deleted and rebuilt in a minute.Tasks 2 and 3
3. The PyTorch buildOne wheel per accelerator: +cu130 (CUDA 13), +rocm7.2, the macOS arm64 wheel with MPS built in, or +cpu. The wheel index you install from decides which.Task 3
4. Driver and runtimeThe NVIDIA driver, at the version its CUDA build needs (Requirements, Track N); ROCm 7.2 user space and the render and video groups; macOS 14 or later for MPS and MLX. Installed by the OS, not by pip.Preflight
5. HardwareGB10 (S), Radeon 8060S (X), the Apple GPU (M), or a GeForce or RTX PRO card (N). An 8 GB machine is enough for this lab.Preflight

pip and uv only ever touch the top three layers. When PyTorch reports no accelerator, the fault is almost always in the two layers below them, which is why the preflight checks those first.

Requirements

The lab is written against PyTorch 2.14.0 · verified 2026-09-12 with torchvision 0.29.0 · verified 2026-09-12, JupyterLab 4.6.3 · verified 2026-09-12, Matplotlib 3.11.2 · verified 2026-09-12 and uv 0.12.11 · verified 2026-09-08, plus MLX 0.32.2 · verified 2026-09-12 on Track M; the Track S container carries NVIDIA’s own builds instead. Every track needs an internet connection and the disk below. Sizes are download sizes read from the package indexes on 2026-09-12; installed size is larger.

Track	PyTorch download	Other downloads	Disk to allow	Attended	Unattended
S, container path	nvcr.io/nvidia/pytorch:26.08-py3, 11.14 GB compressed (NGC catalog)	MNIST 11.6 MB	30 GB	45 min	10–40 min for the pull, then seconds per run
S, wheel path	torch-2.14.0+cu130 aarch64 wheel plus the CUDA library wheels below	jupyterlab, matplotlib, MNIST	8 GB	45 min	5–15 min of downloads, then seconds per run
X	torch-2.14.0+rocm7.2 x86_64 wheel with its ROCm libraries (not tabulated here)	same	8 GB	45 min	5–15 min of downloads; runs take seconds on the GPU, minutes on the CPU
M	torch-2.14.0 macOS arm64 wheel, 127 MB; mlx 0.6 MB plus mlx-metal 42–64 MB depending on the macOS version	same	3 GB	45 min	1–3 min of downloads, then seconds per run
N	torch-2.14.0 Linux x86_64 wheel, 555 MB, plus the CUDA library wheels below	same	8 GB	45 min	5–15 min of downloads, then seconds per run

The CUDA build is large because the CUDA libraries are separate wheels that torch depends on; those with sizes on PyPI for PyTorch 2.14.0 · verified 2026-09-12 on Linux x86_64:

Wheel	Size
torch-2.14.0-cp312-cp312-manylinux_2_28_x86_64.whl	554.6 MB
nvidia_cudnn_cu13-9.24.0.43	553.1 MB
nvidia_nccl_cu13-2.30.7	216.0 MB
nvidia_cusparselt_cu13-0.8.1	170.1 MB
nvidia_nvshmem_cu13-3.4.5	60.4 MB
cuda-toolkit 13.0.3 components (cuBLAS, cuFFT, cuRAND, cuSOLVER, cuSPARSE, NVRTC and others), cuda-bindings, triton 3.8	not tabulated here; several hundred MB more

Memory floor: 8 GB on every track, with no reduced path; the model is 0.8 MB and MNIST is 55 MB decoded.

Track S — NVIDIA DGX Spark

A DGX Spark on DGX OS. The DGX Spark release notes (read 2026-09-12) list DGX OS 7.5.0 with CUDA Toolkit 13.0.2 and GPU driver 580.159.03, and the DGX Spark documentation states that “the NVIDIA Container Toolkit is preinstalled and configured on DGX Spark systems”, so neither path below installs anything at the driver layer:

• Container (primary). The NGC PyTorch container, which the catalog page lists as multi-arch, so the same tag serves the Spark’s aarch64 CPU. The 26.08 release notes list Ubuntu 24.04, Python 3.12, PyTorch 2.14.0a0 (an NVIDIA build), CUDA 13.4.1 and JupyterLab 4.6.3 in the image, and list neither torchvision nor matplotlib, which Task 3 checks for and installs with pip. No uv inside it.
• Wheel (alternative). The PyTorch cu130 index carries manylinux_2_28_aarch64 wheels for PyTorch 2.14.0 · verified 2026-09-12 and torchvision 0.29.0 · verified 2026-09-12, and DGX OS ships the driver CUDA 13 needs; setup-env.sh with TRACK=spark uses that index. Written from the index listing, not validated on a Spark.

Docker needs your user in the docker group, or sudo in front of every docker command; NVIDIA’s page says either is fine.

Track X — AMD Ryzen AI Max+ 395Partial

AMD's PyTorch install page documents a nightly rocm7.2 wheel index and its prerequisites page states that ROCm does not currently support integrated graphics, while the ROCm compatibility matrix lists gfx1151 (both read 2026-09-12). The GPU path is expected to work for this lab's model; the CPU fallback is a supported way to finish, and the page says how.

A Ryzen AI Max+ 395 machine on Ubuntu with the ROCm 7.2 user-space packages installed and your user in the render and video groups (AMD’s prerequisites page gives sudo usermod -a -G video,render $LOGNAME; log out and in afterwards). The wheel’s ROCm version must match the installed one to the minor: the stable PyTorch index has torch-2.14.0+rocm7.2, and AMD’s page documents a nightly rocm7.2 index; setup-env.sh uses the stable one by default and the nightly with ROCM_NIGHTLY=1. On a Windows machine, run Linux for this lab: every step below is written for Ubuntu.

ROCm itself (the amdgpu driver and the ROCm packages, per distribution) comes from AMD’s quick-start installation guide, which Part 5’s Strix Halo page walks through. If ROCm is not installed yet, or the wheel does not see the GPU, run TORCH_BACKEND=cpu TRACK=strix bash setup-env.sh instead: the CPU wheel carries no ROCm libraries and runs the whole lab in a few minutes on the 16-core CPU, a valid result as long as the notebook says so. The ROCm wheel is revisited in Part 5 and Part 11.

Track M — Apple silicon

An Apple silicon Mac on macOS 14 or later, which both the MPS backend notes and the MLX install page require, with a native arm64 Python 3.10 or newer (uv installs one). The PyPI torch wheel for macOS arm64 has MPS built in; there is no separate index. MLX is one more PyPI package and needs nothing else.

Track N — NVIDIA desktop or laptop

A desktop or laptop with an NVIDIA GPU and, for the default wheel, an NVIDIA driver of 580 or later, the CUDA release notes’ minimum for CUDA 13.x applications. On Linux, install the driver from your distribution and nothing else: the wheel brings its own CUDA runtime libraries. With a driver older than 580, setup-env.sh takes TORCH_BACKEND=cu126, which installs torch 2.14.0+cu126 and needs driver 560.28.03 or later on Linux (560.76 on Windows) per the same notes; the cu128 index stops at torch 2.11.0 (index read 2026-09-12), so this course does not use it.

On Windows, this course’s path is WSL2 with Ubuntu, and the CUDA on WSL guide has one firm rule: install the NVIDIA Windows driver, and “do not install any Linux display driver in WSL”, because the Windows driver is mapped into WSL2. Everything else in this lab then happens inside the WSL2 terminal exactly as on Linux. The guide asks for WSL kernel 5.10.16.3 or later; the preflight runs wsl.exe --update and prints the kernel version.

Preflight: check the layers you did not install

Run the block for your track and compare with the expected output; the values go into the notebook’s Machine section in Task 1.

Track S — NVIDIA DGX Spark

RunnableTrack S · DGX Spark

preflight, DGX Spark

nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv
docker --version
groups | tr ' ' '\n' | grep -x docker || echo "not in the docker group: use sudo docker, or add yourself"
df -h ~
uname -m

Output — what you should see

name, driver_version, memory.total [MiB]
NVIDIA GB10, 580.xxx.xx, 1xxxxx MiB
Docker version xx.x.x, build xxxxxxx
docker
Filesystem      Size  Used Avail Use% Mounted on
/dev/nvme0n1p2  9xxG  xxxG  xxxG  xx% /
aarch64

Avail must be at least 30 GB for the container path, 8 GB for the wheel path.

Track X — AMD Ryzen AI Max+ 395Partial

ROCm on gfx1151 is treated as partial; see Requirements.

RunnableTrack X · Ryzen AI Max+

preflight, Ryzen AI Max+ 395

rocminfo | grep -iE "Marketing Name|gfx"
amd-smi version
groups | tr ' ' '\n' | grep -xE "render|video"
df -h ~
nproc

Output — what you should see

  Marketing Name:          <the CPU's name>
  Marketing Name:          <the GPU's name>
  Name:                    gfx1151
AMDSMI Tool: xx.x.x | AMDSMI Library version: xx.x.x | ROCm version: 7.2.x
render
video
Filesystem      Size  Used Avail Use% Mounted on
/dev/nvme0n1p2  9xxG  xxxG  xxxG  xx% /
xx

Two things must be there: a gfx1151 line (the ROCm compatibility matrix’s name for this chip), meaning the runtime enumerates the GPU, and both render and video in your groups. rocminfo and amd-smi version are the checks AMD’s post-installation page gives; the Marketing Name wording is not documented for this chip. No gfx line is a fault at the ROCm layer, below anything Python can fix: take the CPU path from Requirements. Note the ROCm version, which the Task 3 wheel must match to the minor.

Track M — Apple silicon

RunnableTrack M · Apple silicon

preflight, Apple silicon

sw_vers -productVersion
uname -m
system_profiler SPHardwareDataType | grep -E "Chip|Memory"
df -h ~

Output — what you should see

26.x
arm64
      Chip: Apple M4 Pro
      Memory: 48 GB
Filesystem       Size   Used  Avail Capacity iused ifree %iused  Mounted on
/dev/disk3s5    9xxGi  xxxGi  xxxGi    xx%   xxxxx xxxxxx   x%   /System/Volumes/Data

uname -m must print arm64; x86_64 means the terminal runs under Rosetta, and every Python it starts is an x86 Python that neither MPS nor MLX loads into. The macOS version must be 14 or later. This block was not run for this draft; the layout of system_profiler’s output may differ from the sketch above.

Track N — NVIDIA desktop or laptop

On Linux, run only the second block. On Windows, run the PowerShell block, then the second block in the Ubuntu terminal. The PowerShell block is written from Microsoft’s WSL command reference and the CUDA on WSL guide and was not run for this draft.

RunnableTrack N · Windows

preflight, Windows PowerShell

wsl.exe --update
wsl cat /proc/version

Output — what you should see

Checking for updates.
The most recent version of Windows Subsystem for Linux is already installed.   (wording varies by version)
Linux version 6.x.x.x-microsoft-standard-WSL2 (...) ...

RunnableTrack N · NVIDIA GPU

preflight, Linux or the WSL2 terminal

nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv
df -h ~
uname -m

Output — what you should see

name, driver_version, memory.total [MiB]
NVIDIA GeForce RTX 4090, 580.xx.xx, 24564 MiB
Filesystem      Size  Used Avail Use% Mounted on
/dev/nvme0n1p2  9xxG  xxxG  xxxG  xx% /
x86_64

Inside WSL2 the tool lives at /usr/lib/wsl/lib/nvidia-smi, mapped in from the Windows driver; the guide says to use that path or add the directory to PATH if a bare nvidia-smi is not found. driver_version below 580 means using TORCH_BACKEND=cu126 in Task 3 (Requirements, Track N). Avail must be at least 8 GB.

Tasks

1. Make the course directory and start the lab notebook

Everything the labs produce goes in one directory, and the notebook lives at its top. Create the directory now; every later block on this page assumes you are in it.

RunnableAll tracks

create the course directory

mkdir -p ~/llm-course
cd ~/llm-course

Download the notebook template below into that directory as labbook.md (the download link is under the listing), then fill in the Machine section from the preflight: track; machine model, CPU (with the nproc count on X), GPU or chip, memory and storage; OS and version; driver version (S, N), ROCm version (X) or macOS version (M); the date. The Environment section is filled in Task 3.

Fragment — not complete on its own

labbook-template.md

# Lab notebook

One file, kept for the whole course. Every lab ends by adding to it. The capstone is written from it.

## Machine

- Track: (S, X, M or N)
- Machine: (model, CPU, GPU or chip, memory, storage)
- Operating system and version:
- Driver, CUDA, ROCm or macOS version:
- Date prepared:

## Environment

- Python version:
- torch version and build (cuda, rocm, mps):
- mlx version (Track M):
- jupyterlab and matplotlib versions:
- accelerator line printed by setup-env.sh, verbatim:
- Other tools and versions, as they are installed:

## Results

Each lab appends one line per run below, as JSON, or a short table with the same fields:
lab, device, tool versions, the settings that produced the number, the number, the date.

Download labbook-template.md25 lines

One example; the wording does not matter, the fields do:

Fragment — not complete on its own

labbook.md, Machine section, one example

## Machine

- Track: N
- Machine: desktop, Ryzen 9 7950X, GeForce RTX 4090 24 GB, 64 GB RAM, 2 TB NVMe
- Operating system and version: Windows 11 24H2 with WSL2, Ubuntu 24.04
- Driver, CUDA, ROCm or macOS version: NVIDIA driver 580.xx (Windows), WSL kernel 6.x
- Date prepared: 2026-09-12

✦Why a notebook, and why a text file

Every later lab ends by recording a number with the settings that produced it. The scripts append one JSON line per run under Results, so later parts can read them back and the capstone report is written from this file. Back it up like code.

2. Install uv, and know what a virtual environment is

uv creates virtual environments and installs packages, and every lab in this course uses it. Its documentation gives one installer for Linux, macOS and WSL2; Track S needs uv only for the wheel path, but installing it costs nothing.

RunnableAll tracks

install uv (Linux, macOS, WSL2)

curl -LsSf https://astral.sh/uv/install.sh | sh

On a Mac, brew install uv does the same. The installer puts the binary in ~/.local/bin and edits your shell profile to add it to PATH, so open a new terminal, then confirm:

RunnableAll tracks

confirm uv

uv --version
uv python install 3.12

Output — what you should see

uv 0.12.11 (<arch>-<os>)
 Downloaded cpython-3.12.x-<os>-<arch>-<libc> (download)
Installed Python 3.12.x in x.xxs
 + cpython-3.12.x-<os>-<arch>-<libc> (python3.12)

The triple in brackets is the build’s target (x86_64-unknown-linux-gnu on the author’s machine). If a managed 3.12 is already present the second command prints Python 3.12 is already installed instead. It is optional: uv venv --python 3.12 downloads the interpreter by itself if it is missing.

What the environment is, one level down. A virtual environment is a directory. uv venv makes one containing bin/python, which points at a real interpreter that uv manages under ~/.local/share/uv/python/, a lib/python3.12/site-packages/ directory where uv pip install puts packages, and a pyvenv.cfg file that records which interpreter it was made from:

Output — what you should see

home = /home/you/.local/share/uv/python/cpython-3.12-linux-x86_64-gnu/bin
implementation = CPython
uv = 0.12.11
version_info = 3.12
include-system-site-packages = false

(The patch level is recorded only when you ask for one, e.g. --python 3.12.13; the setup script asks for 3.12.)

source .venv/bin/activate does two things that matter: it puts .venv/bin at the front of PATH, so python and jupyter resolve to the environment’s copies, and it sets VIRTUAL_ENV, which is the first place uv looks for the environment to install into. Nothing is installed “into the shell”; a new terminal has neither change, which is why every later page starts with the activate line. include-system-site-packages = false is why the system Python’s packages, and its problems, are invisible from inside.

Command	What it does	Use it when
uv venv --python 3.12 .venv	Creates the directory above; downloads the interpreter if needed	Once per course directory; again with --clear to start over
uv pip install torch torchvision	Installs into the active environment, or .venv in the current directory	Every package in this course
uv pip install --torch-backend=cu130 torch	The same, but resolves PyTorch packages from the named PyTorch index	Choosing the accelerator build; auto queries the installed driver
uv pip list	Lists what is installed, with versions	Filling in the notebook; diagnosing “which torch is this”
uv run python script.py	Runs a command inside the environment without activating it	Scripts started from cron or another shell
pip install ... (bare)	Installs into whatever pip resolves to; a uv environment has no pip unless uv venv --seed was used	Never in this course

3. Create the environment and install PyTorch for your track

The wheel index decides which accelerator your PyTorch can use, and that decision is invisible afterwards except through torch.__version__. PyTorch publishes one wheel per build on its own index at download.pytorch.org/whl/, and PyPI’s torch is one of them chosen per platform:

Track	Where the wheel comes from	torch.__version__ shows	Runtime it needs	The check that proves it
S (wheel)	index cu130, manylinux_2_28_aarch64 (PyPI’s 454 MB aarch64 wheel pulls the same CUDA 13 libraries)	2.14.0+cu130; bare 2.14.0 if it came from PyPI, which is not an error	driver 580 or later (DGX OS 7.5.0 ships 580.159.03)	torch.cuda.is_available() is True; get_device_name(0) is NVIDIA GB10
N	index cu130, chosen by --torch-backend=auto from the installed driver (PyPI’s Linux x86_64 wheel is the same CUDA 13 build); cu126 for older drivers	2.14.0+cu130, or 2.14.0+cu126 with TORCH_BACKEND=cu126; bare 2.14.0 if it came from PyPI	driver 580 or later for CUDA 13; 560.28.03 for cu126	torch.cuda.is_available() is True; torch.version.cuda names the CUDA the wheel was built with
X	index rocm7.2 (stable), or AMD’s documented nightly rocm7.2 index	2.14.0+rocm7.2 or 2.15.0.devYYYYMMDD+rocm7.2	ROCm 7.2 user space, render and video groups	torch.cuda.is_available() is True and torch.version.hip is set; the ROCm build reports the GPU through the cuda device name
M	PyPI, the macosx_14_0_arm64 wheel	2.14.0	macOS 14 or later, arm64 Python	torch.backends.mps.is_available() is True
any, CPU only	index cpu	2.14.0+cpu	nothing	torch.cuda.is_available() is False and stays so

That table is the decision rule for the rest of the course: whenever a framework is slow, print torch.__version__ and the availability check before anything else. setup-env.sh does both for you and encodes the per-track defaults; its variables at the top are the escape hatches.

RunnableAll tracks

setup-env.sh

#!/usr/bin/env bash
# Purpose: create the course's Python environment on one machine: a uv virtual environment
#          holding the PyTorch build for its track (and MLX on a Mac), torchvision, JupyterLab
#          and matplotlib, then print the versions and the accelerator PyTorch can see
# Platform: spark (the aarch64 CUDA wheel; the NGC container is the other Spark path),
#           strix, mac, nvidia (inside WSL2 on Windows)
# Minimum memory: 8 GB
# Assumes: uv is installed and on PATH; internet access to download wheels; run from the
#          course directory (~/llm-course), which will hold .venv
#
# Usage: TRACK=<spark|strix|mac|nvidia> bash setup-env.sh
#   Optional environment variables:
#     PYTHON=3.12            Python version for the environment (uv downloads it if missing)
#     VENV=.venv             where to create it
#     FRESH=1                replace an existing .venv instead of reusing it
#     TORCH_BACKEND=<value>  override the track's PyTorch index: cu130, cu126, rocm7.2,
#                            cpu or auto (values uv pip install --torch-backend accepts;
#                            cu130 needs an NVIDIA driver of 580 or later, cu126 one of
#                            560.28.03 or later; cu128 resolves to torch 2.11.0, not 2.14.0,
#                            so this course does not use it)
#     ROCM_NIGHTLY=1         Track X only: use AMD's documented nightly index with --pre
#                            instead of the stable rocm7.2 index
#
# Track defaults: spark -> the cu130 index (has aarch64 wheels); nvidia -> auto (uv queries
# the installed driver and picks the most compatible CUDA index); strix -> rocm7.2;
# mac -> the PyPI wheel, which is the MPS build on Apple silicon, plus mlx.
set -euo pipefail

TRACK="${TRACK:-}"
PYTHON="${PYTHON:-3.12}"
VENV="${VENV:-.venv}"
FRESH="${FRESH:-0}"
TORCH_BACKEND="${TORCH_BACKEND:-}"
ROCM_NIGHTLY="${ROCM_NIGHTLY:-0}"

die() { echo "setup-env: $*" >&2; exit 1; }
command -v uv >/dev/null || die "uv is not installed; see https://docs.astral.sh/uv/getting-started/installation/ and rerun"
case "$TRACK" in
  spark|strix|mac|nvidia) ;;
  "") die "set TRACK=spark, TRACK=strix, TRACK=mac or TRACK=nvidia" ;;
  *) die "unknown TRACK '$TRACK' (expected spark, strix, mac or nvidia)" ;;
esac
if [[ "$TRACK" == "mac" && "$(uname -m)" != "arm64" ]]; then
  die "TRACK=mac needs an Apple silicon Mac; uname -m printed $(uname -m)"
fi

echo "==> uv $(uv --version | awk '{print $2}')"
if [[ -f "$VENV/pyvenv.cfg" && "$FRESH" != "1" ]]; then
  echo "==> Reusing the existing environment at $VENV (set FRESH=1 to replace it)"
elif [[ -f "$VENV/pyvenv.cfg" ]]; then
  echo "==> Replacing the existing environment at $VENV"
  uv venv --clear --python "$PYTHON" "$VENV"
else
  echo "==> Creating $VENV with Python $PYTHON"
  uv venv --python "$PYTHON" "$VENV"
fi
# shellcheck disable=SC1091
source "$VENV/bin/activate"

case "$TRACK" in
  spark)
    BACKEND="${TORCH_BACKEND:-cu130}"
    echo "==> Installing PyTorch (CUDA build, aarch64 wheel) and torchvision from the $BACKEND index"
    uv pip install torch torchvision --torch-backend="$BACKEND"
    ;;
  nvidia)
    BACKEND="${TORCH_BACKEND:-auto}"
    echo "==> Installing PyTorch (CUDA build) and torchvision with --torch-backend=$BACKEND"
    uv pip install torch torchvision --torch-backend="$BACKEND"
    ;;
  strix)
    if [[ "$ROCM_NIGHTLY" == "1" ]]; then
      echo "==> Installing PyTorch (ROCm build) and torchvision from AMD's documented nightly index"
      uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm7.2
    else
      BACKEND="${TORCH_BACKEND:-rocm7.2}"
      echo "==> Installing PyTorch (ROCm build) and torchvision with --torch-backend=$BACKEND"
      uv pip install torch torchvision --torch-backend="$BACKEND"
    fi
    ;;
  mac)
    echo "==> Installing PyTorch (the macOS arm64 wheel; MPS is built in), torchvision and MLX"
    uv pip install torch torchvision mlx
    ;;
esac

echo "==> Installing JupyterLab and matplotlib"
uv pip install jupyterlab matplotlib

echo "==> Verifying"
python - <<'PY'
import platform
import torch
print(f"python {platform.python_version()} ({platform.machine()})")
print(f"torch {torch.__version__}   cuda build: {torch.version.cuda}   hip build: {torch.version.hip}")
mps = getattr(torch.backends, "mps", None)
if torch.cuda.is_available():
    print(f"accelerator: cuda / {torch.cuda.get_device_name(0)}")
elif mps is not None and mps.is_available():
    print("accelerator: mps (Apple silicon GPU)")
else:
    built = "built" if (mps is not None and mps.is_built()) else "not built"
    print(f"accelerator: none visible; the lab will run on the CPU (mps support {built})")
PY
if [[ "$TRACK" == "mac" ]]; then
  python -c 'import mlx.core as mx; print("mlx", mx.__version__, "default device:", mx.default_device())'
fi
echo "jupyterlab $(jupyter lab --version)"
python -c 'import matplotlib; print("matplotlib", matplotlib.__version__)'
echo "==> Done. Copy the lines above into the Environment section of labbook.md."
echo "    Activate in every new terminal with: source $VENV/bin/activate"

Download setup-env.sh111 lines

Track S — NVIDIA DGX Spark

Container path. Look up the newest release tag on the container’s NGC catalog page (tags have the form yy.mm-py3; on 2026-09-12 it was 26.08-py3), then run the script with it. It pulls the image, checks that PyTorch inside it can see the GPU, and drops you into a shell with your course directory mounted at /workspace/course and port 8888 published on 127.0.0.1 only, for JupyterLab in Task 4.

RunnableTrack S · DGX Spark

setup-env-spark.sh

#!/usr/bin/env bash
# Purpose: pull the NVIDIA PyTorch container on a DGX Spark, check that PyTorch inside it sees
#          the GPU, then start an interactive shell in it with the course directory mounted
#          at /workspace/course and port 8888 published for JupyterLab
# Platform: spark
# Minimum memory: 8 GB
# Assumes: DGX OS with Docker and the NVIDIA Container Toolkit (preinstalled and configured on
#          DGX Spark per NVIDIA's documentation); your user in the docker group, or run with
#          sudo; about 12 GB of free disk for the image; run from the course directory
#          (~/llm-course)
#
# Usage: TAG=<yy.mm-py3> bash setup-env-spark.sh
#   TAG is a release tag from the container's NGC catalog page (26.08-py3 was the newest on
#   2026-09-12). The image is multi-arch, so the same tag serves the Spark's aarch64 CPU.
#   PyTorch and JupyterLab are in the image (NGC catalog page); the 26.08 release notes list
#   neither torchvision nor matplotlib, so the GPU check below also reports whether each
#   is present.
set -euo pipefail

TAG="${TAG:?set TAG to a release tag from the NGC catalog page, of the form yy.mm-py3}"
IMAGE="nvcr.io/nvidia/pytorch:${TAG}"
PORT="${PORT:-8888}"

command -v docker >/dev/null || { echo "setup-env-spark: docker is not installed" >&2; exit 1; }
if ! command -v nvidia-smi >/dev/null; then
  echo "setup-env-spark: nvidia-smi is not on PATH on the host; the GPU check below will fail" >&2
fi

echo "==> Pulling $IMAGE (about 11 GB compressed; the first pull takes a while)"
docker pull "$IMAGE"

echo "==> Checking that the GPU is visible to PyTorch inside the container"
docker run --gpus all --rm --interactive "$IMAGE" python - <<'PY'
import importlib.util
import torch
print("torch", torch.__version__, "cuda build:", torch.version.cuda)
if torch.cuda.is_available():
    print("accelerator: cuda /", torch.cuda.get_device_name(0))
else:
    print("accelerator: none visible")
for m in ("torchvision", "matplotlib"):
    print(f"{m}: " + ("present" if importlib.util.find_spec(m) else "missing; run: pip install " + m))
PY

echo "==> Starting an interactive shell with $PWD mounted at /workspace/course"
echo "    Port $PORT is published on 127.0.0.1 only; inside the container JupyterLab starts with:"
echo "      jupyter lab --ip 0.0.0.0 --port $PORT --no-browser --allow-root"
exec docker run --gpus all --interactive --tty --rm --ipc=host \
  --publish "127.0.0.1:${PORT}:${PORT}" \
  --volume "$PWD:/workspace/course" --workdir /workspace/course \
  "$IMAGE" bash

Download setup-env-spark.sh51 lines

RunnableTrack S · DGX Spark

pull the container and start a shell in it

cd ~/llm-course
TAG=26.08-py3 bash setup-env-spark.sh

Output — what you should see

==> Pulling nvcr.io/nvidia/pytorch:26.08-py3 (about 11 GB compressed; the first pull takes a while)
26.08-py3: Pulling from nvidia/pytorch
...
Status: Downloaded newer image for nvcr.io/nvidia/pytorch:26.08-py3
==> Checking that the GPU is visible to PyTorch inside the container
torch 2.14.0a0+4fdf77b940 cuda build: 13.x
accelerator: cuda / NVIDIA GB10
torchvision: present
matplotlib: present
==> Starting an interactive shell with /home/you/llm-course mounted at /workspace/course
    Port 8888 is published on 127.0.0.1 only; inside the container JupyterLab starts with:
      jupyter lab --ip 0.0.0.0 --port 8888 --no-browser --allow-root
root@<container-id>:/workspace/course#

The torch version inside the container is NVIDIA’s build (2.14.0a0+... for tag 26.08); record it as printed. If either of the last two lines says missing, run pip install torchvision matplotlib in the container shell before Task 5 and note the versions. Stay in that shell for the rest of the lab: files the scripts write appear in ~/llm-course on the host, and python is the container’s own, so skip the “activate the environment” line below. If the check prints accelerator: none visible, see Troubleshooting.

Wheel path, if you prefer to work outside a container. The same setup script as the other tracks, with TRACK=spark, which installs from the cu130 index:

RunnableTrack S · DGX Spark

wheel path: environment on the host

cd ~/llm-course
TRACK=spark bash setup-env.sh

Expected output is the Track N block below with NVIDIA GB10 as the device name and aarch64 as the machine. Unvalidated on hardware at the time of writing; if accelerator: says none, go back to the container and note in the notebook which path you used.

Track X — AMD Ryzen AI Max+ 395Partial

Stable rocm7.2 wheel by default, AMD's nightly index with ROCM_NIGHTLY=1; the CPU fallback is a supported way to finish.

RunnableTrack X · Ryzen AI Max+

environment with the ROCm build

cd ~/llm-course
TRACK=strix bash setup-env.sh

Output — what you should see

==> uv 0.12.11
==> Creating .venv with Python 3.12
Using CPython 3.12.x
Creating virtual environment at: .venv
Activate with: source .venv/bin/activate
==> Installing PyTorch (ROCm build) and torchvision with --torch-backend=rocm7.2
Resolved xx packages in x.xxs
...
 + torch==2.14.0+rocm7.2
 + torchvision==0.29.0+rocm7.2
==> Installing JupyterLab and matplotlib
...
==> Verifying
python 3.12.x (x86_64)
torch 2.14.0+rocm7.2   cuda build: None   hip build: 7.2.xxxxx
accelerator: cuda / AMD Radeon Graphics
jupyterlab 4.x.x
matplotlib 3.x.x
==> Done. Copy the lines above into the Environment section of labbook.md.
    Activate in every new terminal with: source .venv/bin/activate

accelerator: cuda / AMD Radeon ... is correct, as the table above explains; hip build: carrying a version is what tells a ROCm wheel from a CUDA one. The device-name string for this chip is not documented and may differ. If it prints accelerator: none visible, the lab still runs on the CPU, which the training script picks automatically; the Track X Troubleshooting row has the three things to try first, including ROCM_NIGHTLY=1 TRACK=strix bash setup-env.sh for AMD’s documented nightly index.

Track M — Apple silicon

RunnableTrack M · Apple silicon

environment with PyTorch (MPS) and MLX

cd ~/llm-course
TRACK=mac bash setup-env.sh

Output — what you should see

==> uv 0.12.11
==> Creating .venv with Python 3.12
Using CPython 3.12.x
Creating virtual environment at: .venv
Activate with: source .venv/bin/activate
==> Installing PyTorch (the macOS arm64 wheel; MPS is built in), torchvision and MLX
Resolved xx packages in x.xxs
...
 + mlx==0.32.x
 + mlx-metal==0.32.x
 + torch==2.14.0
 + torchvision==0.29.0
==> Installing JupyterLab and matplotlib
...
==> Verifying
python 3.12.x (arm64)
torch 2.14.0   cuda build: None   hip build: None
accelerator: mps (Apple silicon GPU)
mlx 0.32.x default device: Device(gpu, 0)
jupyterlab 4.x.x
matplotlib 3.x.x
==> Done. Copy the lines above into the Environment section of labbook.md.
    Activate in every new terminal with: source .venv/bin/activate

Three lines matter: accelerator: mps from PyTorch, default device: Device(gpu, 0) from MLX, and (arm64) after the Python version; (x86_64) there is the Rosetta case from the preflight.

Track N — NVIDIA desktop or laptop

RunnableTrack N · NVIDIA GPU

environment with the CUDA build (Linux, or the WSL2 terminal)

cd ~/llm-course
TRACK=nvidia bash setup-env.sh

Output — what you should see

==> uv 0.12.11
==> Creating .venv with Python 3.12
Using CPython 3.12.x
Creating virtual environment at: .venv
Activate with: source .venv/bin/activate
==> Installing PyTorch (CUDA build) and torchvision with --torch-backend=auto
Resolved xx packages in x.xxs
...
 + nvidia-cudnn-cu13==9.24.0.43
 + torch==2.14.0+cu130
 + torchvision==0.29.0+cu130
==> Installing JupyterLab and matplotlib
...
==> Verifying
python 3.12.x (x86_64)
torch 2.14.0+cu130   cuda build: 13.0   hip build: None
accelerator: cuda / NVIDIA GeForce RTX 4090
jupyterlab 4.x.x
matplotlib 3.x.x
==> Done. Copy the lines above into the Environment section of labbook.md.
    Activate in every new terminal with: source .venv/bin/activate

--torch-backend=auto is uv’s documented selector: it queries the installed CUDA driver, picks the most compatible PyTorch index, and falls back to the CPU build when it finds no GPU. Check the torch line. +cpu or no +cu... suffix means the driver was not detected: rerun with TORCH_BACKEND=cu130 FRESH=1 TRACK=nvidia bash setup-env.sh. A version other than 2.14.0 (for example 2.11.0+cu128) means auto matched an older driver to an older index: rerun with TORCH_BACKEND=cu126 FRESH=1 TRACK=nvidia bash setup-env.sh, or update the driver to 580. accelerator: none visible with a working nvidia-smi is in Troubleshooting.

Whatever the script printed, activate the environment in every new terminal before running anything else, and prove that python is the environment’s:

RunnableAll tracks

activate the environment, then check which python this is

cd ~/llm-course
source .venv/bin/activate
which python
python -c "import torch; print(torch.__version__)"

Output — what you should see

/home/you/llm-course/.venv/bin/python
2.14.0+cu130

Record in the Environment section of the notebook: Python version and machine, torch version with its suffix, cuda build or hip build, the accelerator: line verbatim, mlx version and default device (M), jupyterlab and matplotlib versions, and, on Track S, whether you are in the container or on the wheel. If which python prints /usr/bin/python or python: not found, the environment is not active; see Troubleshooting.

4. Start JupyterLab and make the course notebook file

JupyterLab, installed by the setup script and already in the NGC container, serves a web page from your course directory in which code runs in cells against the same environment. The course’s rule, which Part 11 repeats: scripts for anything you will record, notebooks for looking. A notebook’s cells can run in any order, so a result in one is not reproducible from the file; a script with arguments is.

Track S — NVIDIA DGX Spark

Inside the container shell from Task 3, where port 8888 is already published to 127.0.0.1, start the server on all container interfaces and as root, the container’s user:

RunnableTrack S · DGX Spark

JupyterLab inside the container

jupyter lab --ip 0.0.0.0 --port 8888 --no-browser --allow-root

Open the http://127.0.0.1:8888/lab?token=... line it prints in a browser on the Spark’s desktop. From another machine, tunnel first: ssh -L 8888:127.0.0.1:8888 you@spark in a terminal on that machine, then open the same URL there.

Track X — AMD Ryzen AI Max+ 395

RunnableTrack X · Ryzen AI Max+

JupyterLab

cd ~/llm-course
source .venv/bin/activate
jupyter lab

The browser opens by itself on the desktop. Over SSH, add --no-browser, tunnel with ssh -L 8888:127.0.0.1:8888 you@machine from your desktop, and open the printed URL there.

Track M — Apple silicon

RunnableTrack M · Apple silicon

JupyterLab

cd ~/llm-course
source .venv/bin/activate
jupyter lab

The browser opens by itself.

Track N — NVIDIA desktop or laptop

RunnableTrack N · NVIDIA GPU

JupyterLab (Linux, or the WSL2 terminal)

cd ~/llm-course
source .venv/bin/activate
jupyter lab --no-browser

Inside WSL2 there may be no browser for the server to open, so --no-browser is used and you paste the URL into a Windows browser; localhost inside WSL2 is reachable from Windows as localhost. On a Linux desktop you can drop --no-browser.

The terminal shows the server’s log; the lines to look for, captured from JupyterLab 4.6.3:

Output — what you should see

[I 2026-09-12 11:49:31.600 ServerApp] Serving notebooks from local directory: /home/you/llm-course
[I 2026-09-12 11:49:31.600 ServerApp] Jupyter Server 2.21.0 is running at:
[I 2026-09-12 11:49:31.600 ServerApp] http://localhost:8888/lab?token=<48 hex characters>
[I 2026-09-12 11:49:31.600 ServerApp]     http://127.0.0.1:8888/lab?token=<48 hex characters>
[I 2026-09-12 11:49:31.600 ServerApp] Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).

Serving notebooks from local directory must name your course directory: the notebook file is created there, and %run in a cell finds the scripts by relative path. The token in the URL is the only authentication; the server listens on localhost only, which is where it stays for the whole course.

⛨Security

A Jupyter kernel is a shell with your permissions. Never start JupyterLab with --ip 0.0.0.0 on a machine others can reach; the Track S container does so only inside its own network namespace, which Docker publishes to 127.0.0.1 on the host.

In the browser, choose File → New → Notebook, pick the Python 3 (ipykernel) kernel, and save it as part-01.ipynb (File → Rename Notebook…). Put this in the first cell and run it with Shift+Enter; it is the accelerator check from Task 3, now from inside the kernel, which proves the kernel is the environment you built and not some other Python:

Fragment — not complete on its own

part-01.ipynb, cell 1

import sys, torch
print(sys.executable)
print(torch.__version__)
print("cuda:", torch.cuda.is_available(), " mps:", torch.backends.mps.is_available())

Output — what you should see

/home/you/llm-course/.venv/bin/python
2.14.0+cu130
cuda: True  mps: False

In the NGC container the first line is the image’s own Python under /usr/; the version line and cuda: True are what matter there. Leave the notebook open; cells 2 and 3 are written in Tasks 6 and 9. Stop the server later with Control-C twice in its terminal. Nothing to record: the notebook file is the artefact.

5. Train the model and read every line it prints

Download the training script into the course directory and read it before running it. Every line of the loop from the first lesson is marked with a comment: forward, loss, backward, step.

RunnableAll tracks

train-mnist.py

"""Train a two-layer network on MNIST and report training and validation loss per epoch.

Purpose: the first training run of the course: watch the loss fall, see the validation
         curve, save the weights from the best epoch, and record the run in the lab notebook.
Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, or the CPU);
          the device is chosen automatically and printed, or forced with --device.
Minimum memory: 8 GB
Assumes: torch and torchvision are installed in the active environment, and about
         70 MB of free disk under ./data for the MNIST download (11.6 MB compressed).

Usage: python train-mnist.py [--epochs 5] [--lr 0.1] [--batch-size 128] [--hidden 256]
                             [--train-size 50000] [--device auto|cpu|cuda|mps]
                             [--checkpoint mnist_mlp.pt] [--label first-run]
                             [--labbook labbook.md] [--seed 0] [--data ./data]

Every epoch prints one line: mean training loss over the epoch, validation loss and
accuracy at the end of the epoch, the seconds the epoch took, and "(saved)" when the
validation loss is the lowest seen so far and the weights were written to --checkpoint.
The test set is evaluated once, with the saved checkpoint, after the last epoch.
"""
import argparse
import json
import math
import sys
import time
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, random_split
from torchvision import datasets, transforms


def pick_device(name: str) -> torch.device:
    """Return the requested device, or the best available one for "auto"."""
    mps = getattr(torch.backends, "mps", None)
    mps_ok = mps is not None and mps.is_available()
    if name == "auto":
        if torch.cuda.is_available():
            return torch.device("cuda")
        if mps_ok:
            return torch.device("mps")
        return torch.device("cpu")
    if name == "cuda" and not torch.cuda.is_available():
        sys.exit("train-mnist: --device cuda requested but torch.cuda.is_available() is False")
    if name == "mps" and not mps_ok:
        sys.exit("train-mnist: --device mps requested but torch.backends.mps.is_available() is False")
    return torch.device(name)


def describe_device(device: torch.device) -> str:
    if device.type == "cuda":
        return f"cuda / {torch.cuda.get_device_name(0)}"
    if device.type == "mps":
        return "mps (Apple silicon GPU)"
    return "cpu"


class TwoLayerNet(nn.Module):
    """784 inputs -> hidden units with ReLU -> 10 outputs, one per digit."""

    def __init__(self, hidden: int = 256) -> None:
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(nn.Linear(28 * 28, hidden), nn.ReLU(), nn.Linear(hidden, 10))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(self.flatten(x))


def evaluate(model: nn.Module, loader: DataLoader, device: torch.device, loss_fn: nn.Module) -> tuple[float, float]:
    """Average loss and accuracy over a loader, with gradients switched off."""
    model.eval()
    total_loss, correct, seen = 0.0, 0, 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            logits = model(images)
            total_loss += loss_fn(logits, labels).item() * labels.size(0)
            correct += (logits.argmax(dim=1) == labels).sum().item()
            seen += labels.size(0)
    return total_loss / seen, correct / seen


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--lr", type=float, default=0.1, help="learning rate for plain SGD")
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--hidden", type=int, default=256, help="width of the hidden layer")
    parser.add_argument("--train-size", type=int, default=50_000,
                        help="how many of the 50,000 training images to train on (fewer overfits sooner)")
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps"])
    parser.add_argument("--checkpoint", default="mnist_mlp.pt", help="where the best weights are saved")
    parser.add_argument("--label", default="", help="a short name for this run in the notebook")
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--data", default="./data")
    args = parser.parse_args()
    if not 1 <= args.train_size <= 50_000:
        sys.exit("train-mnist: --train-size must be between 1 and 50000")

    torch.manual_seed(args.seed)
    device = pick_device(args.device)
    print(f"torch {torch.__version__}")
    print(f"device: {describe_device(device)}")

    # MNIST ships as 60,000 training and 10,000 test images. The test set is spent once,
    # at the end; 10,000 of the training images are held out as the validation set.
    to_tensor = transforms.ToTensor()  # 28x28 bytes 0..255 -> float32 tensor (1, 28, 28) in 0.0..1.0
    full_train = datasets.MNIST(args.data, train=True, download=True, transform=to_tensor)
    test_set = datasets.MNIST(args.data, train=False, download=True, transform=to_tensor)
    train_set, val_set = random_split(full_train, [50_000, 10_000], generator=torch.Generator().manual_seed(args.seed))
    if args.train_size < 50_000:
        train_set = Subset(train_set, range(args.train_size))
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_set, batch_size=1000)
    test_loader = DataLoader(test_set, batch_size=1000)
    steps = math.ceil(len(train_set) / args.batch_size)
    print(f"data: train {len(train_set):,}  val {len(val_set):,}  test {len(test_set):,}  "
          f"batch {args.batch_size} -> {steps} steps per epoch")

    model = TwoLayerNet(args.hidden).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"parameters: {n_params:,}")
    for name, p in model.named_parameters():
        print(f"  {name:20s} shape {tuple(p.shape)}")

    loss_fn = nn.CrossEntropyLoss()
    optimiser = torch.optim.SGD(model.parameters(), lr=args.lr)

    # The untrained network: guessing among ten classes costs ln(10) = 2.3026 per image.
    val_loss, val_acc = evaluate(model, val_loader, device, loss_fn)
    print(f"epoch  0  (untrained)          val loss {val_loss:.4f}  val acc {val_acc:.4f}")

    history = []
    best_val, best_epoch = float("inf"), 0
    started = time.time()
    for epoch in range(1, args.epochs + 1):
        epoch_started = time.time()
        model.train()
        running, seen = 0.0, 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            logits = model(images)          # forward
            loss = loss_fn(logits, labels)  # loss
            optimiser.zero_grad()
            loss.backward()                 # backward
            optimiser.step()                # step
            running += loss.item() * labels.size(0)
            seen += labels.size(0)
        train_loss = running / seen
        val_loss, val_acc = evaluate(model, val_loader, device, loss_fn)
        seconds = time.time() - epoch_started
        history.append({"epoch": epoch, "train_loss": round(train_loss, 4), "val_loss": round(val_loss, 4),
                        "val_acc": round(val_acc, 4), "seconds": round(seconds, 1)})
        marker = ""
        if val_loss < best_val:  # False when val_loss is nan, so a diverged epoch is never saved
            best_val, best_epoch = val_loss, epoch
            torch.save(model.state_dict(), args.checkpoint)
            marker = "  (saved)"
        print(f"epoch {epoch:2d}  train loss {train_loss:.4f}  val loss {val_loss:.4f}  "
              f"val acc {val_acc:.4f}  {seconds:5.1f} s{marker}")
    elapsed = time.time() - started

    # The test set is looked at once, with the checkpoint from the validation minimum.
    if best_epoch == 0:
        test_loss, test_acc = float("nan"), float("nan")
        print(f"no epoch improved the validation loss: nothing saved to {args.checkpoint}, "
              f"test set not used  ({elapsed:.0f} s)")
    else:
        model.load_state_dict(torch.load(args.checkpoint, map_location=device))
        test_loss, test_acc = evaluate(model, test_loader, device, loss_fn)
        print(f"best epoch {best_epoch}: test loss {test_loss:.4f}  test acc {test_acc:.4f}  "
              f"({elapsed:.0f} s total, checkpoint {args.checkpoint})")

    if args.labbook:
        record = {
            "lab": "part-01/train-mnist", "label": args.label, "device": describe_device(device),
            "torch": torch.__version__, "epochs": args.epochs, "lr": args.lr, "batch_size": args.batch_size,
            "hidden": args.hidden, "train_size": len(train_set), "parameters": n_params,
            "best_epoch": best_epoch, "test_acc": None if math.isnan(test_acc) else round(test_acc, 4),
            "seconds": round(elapsed, 1), "checkpoint": args.checkpoint, "history": history,
        }
        with Path(args.labbook).open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record) + "\n")
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()

Download train-mnist.py191 lines

What the pieces are, in the order the script meets them:

• Data. datasets.MNIST(..., download=True) fetches four files into ./data/MNIST/raw/ and decodes them; ToTensor() turns each 28×28 image of bytes 0–255 into a float32 tensor of shape (1, 28, 28) with values 0.0–1.0; nn.Flatten() in the model makes it a vector of 784. random_split with a seeded generator holds out 10,000 of the 60,000 training images as the validation set, so the 10,000 test images are touched exactly once, at the end.
• Model. nn.Linear(784, 256) then ReLU then nn.Linear(256, 10): the two-layer network from the neural-networks lesson. Ten outputs, one score (a logit) per digit.
• Loss. nn.CrossEntropyLoss turns the ten logits into probabilities with a softmax and charges -log(probability of the correct digit). A network that spreads probability evenly over ten digits pays -log(1/10) = ln(10) = 2.3026 per image; that is the number an untrained network prints, and the number a broken one gets stuck at.
• Optimiser. Plain SGD: after loss.backward() has put a gradient on every parameter, optimiser.step() does w = w - lr * grad for each of the 203,530 of them.
• Batches and epochs. 50,000 images in batches of 128 is ceil(50000 / 128) = 391 steps per epoch, the last batch holding 80 images. Five epochs are 1,955 steps.

Run it with the defaults, a label and the notebook flag:

RunnableAll tracks

first training run

python train-mnist.py --epochs 5 --label first-run --checkpoint mnist-first.pt --labbook labbook.md

It downloads MNIST first. Below is the author’s run on a CPU build (torch 2.14.0+cpu, seed 0, a shared 12-core x86 machine); your loss digits will differ in the third or fourth decimal, and the seconds column is your machine’s:

Output — what you should see

first-run, author's CPU run

torch 2.14.0+cpu
device: cpu
data: train 50,000  val 10,000  test 10,000  batch 128 -> 391 steps per epoch
parameters: 203,530
  layers.0.weight      shape (256, 784)
  layers.0.bias        shape (256,)
  layers.2.weight      shape (10, 256)
  layers.2.bias        shape (10,)
epoch  0  (untrained)          val loss 2.3091  val acc 0.0645
epoch  1  train loss 0.6121  val loss 0.3412  val acc 0.9044    7.0 s  (saved)
epoch  2  train loss 0.3040  val loss 0.2804  val acc 0.9206   50.7 s  (saved)
epoch  3  train loss 0.2518  val loss 0.2460  val acc 0.9317   40.1 s  (saved)
epoch  4  train loss 0.2164  val loss 0.2086  val acc 0.9423   53.0 s  (saved)
epoch  5  train loss 0.1883  val loss 0.1937  val acc 0.9441   71.5 s  (saved)
best epoch 5: test loss 0.1790  test acc 0.9465  (222 s total, checkpoint mnist-first.pt)
recorded in labbook.md

Read the lines against the arithmetic:

Line	Where the number comes from
parameters: 203,530	784 × 256 + 256 + 256 × 10 + 10; the four shapes below it are those four terms
epoch 0 ... val loss 2.3091 val acc 0.0645	ln(10) = 2.3026, the cost of guessing evenly, plus whatever the random initial weights prefer; a ten-sided coin scores 0.10 on average, and these weights happened to favour the wrong digits
epoch 1 ... train loss 0.6121 val loss 0.3412	Training loss is the mean over the epoch while the weights were still moving down from 2.31; validation loss is measured once, at the end, with the finished weights, so it is lower in epoch 1. Not a bug, and not overfitting in reverse
(saved) on every epoch	Validation loss set a new minimum each time, so the checkpoint was rewritten each time; five epochs is too few for a model this small to overfit 50,000 images
best epoch 5: test acc 0.9465	The test set, seen once, with the epoch-5 weights loaded back from mnist-first.pt
7.0 s then 50.7 s per epoch	The machine, not the model: this run shared its CPU with other jobs from epoch 2 on. On an accelerator, well under a second per epoch is normal

=Expected result

Both losses fall every epoch, and test accuracy after five epochs of SGD at learning rate 0.1 lands around 0.94 to 0.95. Below 0.90, a loss that is not falling, or nan: see Troubleshooting before going on.

What to record: the JSON line is already in labbook.md. Add by hand, next to it or in the Machine section, the device string and the seconds for five epochs, so that later parts can compare the same script on other machines.

6. Plot and read the curves

The columns are the curves. The script below reads every part-01/train-mnist record in the notebook and draws training loss (solid), validation loss (dashed) and a dot on the epoch whose checkpoint was kept, one panel per run, on a logarithmic axis so that a run that reached 0.01 and a run stuck at 2.3 are both readable.

RunnableAll tracks

plot-curves.py

"""Plot the training and validation loss curves recorded in the lab notebook.

Purpose: turn the per-epoch numbers that train-mnist.py and train-mnist-mlx.py appended to
         labbook.md into one picture per run, so the shapes from the generalisation lesson
         can be seen rather than read off a column of digits.
Platform: all (no accelerator needed; matplotlib only)
Minimum memory: 8 GB
Assumes: matplotlib is installed in the active environment, and labbook.md holds at least
         one JSON line whose "lab" field starts with "part-01/train-mnist".

Usage: python plot-curves.py [--labbook labbook.md] [--out curves.png] [--last N]

Each run becomes one panel: training loss (solid), validation loss (dashed) and a dot on
the epoch whose checkpoint was kept. The y axis is logarithmic so a run that reached 0.01
and a run stuck at 2.3 are both readable.
"""
import argparse
import json
import sys
from pathlib import Path

try:
    import matplotlib.pyplot as plt
except ImportError:
    sys.exit("plot-curves: matplotlib is not installed; run: uv pip install matplotlib "
             "(pip install matplotlib inside the NGC container)")


def load_runs(labbook: Path) -> list[dict]:
    runs = []
    for line in labbook.read_text(encoding="utf-8").splitlines():
        if not line.startswith("{"):
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if str(record.get("lab", "")).startswith("part-01/train-mnist") and record.get("history"):
            runs.append(record)
    return runs


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--out", default="curves.png")
    parser.add_argument("--last", type=int, default=0, help="plot only the last N runs (0 = all)")
    args = parser.parse_args()

    labbook = Path(args.labbook)
    if not labbook.exists():
        sys.exit(f"plot-curves: {labbook} does not exist; run train-mnist.py with --labbook first")
    runs = load_runs(labbook)
    if not runs:
        sys.exit(f"plot-curves: no part-01/train-mnist records with a history in {labbook}")
    if args.last:
        runs = runs[-args.last:]

    cols = 2 if len(runs) > 1 else 1
    rows = (len(runs) + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(5.5 * cols, 3.6 * rows), squeeze=False)
    for ax in axes.flat[len(runs):]:
        ax.set_visible(False)

    print(f"{'run':<4} {'label':<16} {'lr':>8} {'epochs':>6} {'train':>6} {'best':>4} {'min val loss':>12} {'final train':>11}")
    for i, (run, ax) in enumerate(zip(runs, axes.flat), start=1):
        epochs = [h["epoch"] for h in run["history"]]
        train = [h["train_loss"] for h in run["history"]]
        val = [h["val_loss"] for h in run["history"]]
        best = run.get("best_epoch") or 0
        ax.plot(epochs, train, marker=".", label="training loss")
        ax.plot(epochs, val, marker=".", linestyle="--", label="validation loss")
        if best:
            ax.plot([best], [val[best - 1]], marker="o", markersize=9, linestyle="none",
                    label=f"kept checkpoint (epoch {best})")
        ax.set_yscale("log")
        ax.set_xlabel("epoch")
        ax.set_ylabel("cross-entropy loss")
        label = run.get("label") or f"run {i}"
        ax.set_title(f"{label}: lr {run['lr']}, {run.get('train_size', 50000):,} images, {run['device']}", fontsize=9)
        ax.grid(True, which="both", alpha=0.3)
        ax.legend(fontsize=8)
        min_val = min(v for v in val if v == v) if any(v == v for v in val) else float("nan")
        print(f"{i:<4} {label[:16]:<16} {run['lr']:>8} {run['epochs']:>6} {run.get('train_size', 50000):>6} "
              f"{best:>4} {min_val:>12.4f} {train[-1]:>11.4f}")

    fig.tight_layout()
    fig.savefig(args.out, dpi=120)
    print(f"wrote {args.out} ({len(runs)} run(s) from {labbook})")


if __name__ == "__main__":
    main()

Download plot-curves.py93 lines

RunnableAll tracks

plot every run recorded so far

python plot-curves.py --labbook labbook.md --out curves.png

Output — what you should see

plot-curves.py after the first run, author's run

run  label                  lr epochs  train best min val loss final train
1    first-run             0.1      5  50000    5       0.1937      0.1883
wrote curves.png (1 run(s) from labbook.md)

Open curves.png, or run the same thing inside the notebook, where the figure appears under the cell:

Fragment — not complete on its own

part-01.ipynb, cell 2

%run plot-curves.py --labbook labbook.md --out curves.png

With one run there is little to read yet: two lines going down. The shapes worth recognising, from the generalisation lesson, and what to do about each:

Shape	What it means	What to do
Both falling, validation slightly above training	Learning; not yet overfitting	Keep training, or stop when the epochs cost more than the gain
Training keeps falling, validation flattens then rises	Overfitting from the turn onwards; the model is fitting the training set’s noise	Keep the checkpoint from the minimum (the script already does); more data, a smaller model or regularisation to move the turn later
Both flat near 2.30	Nothing learned: learning rate too small, a dead network, or labels shuffled against images	Check the learning rate first, then the data pipeline
Training loss enormous in epoch 1, then flat at or above 2.30, or nan	Learning rate too large; the steps overshoot, the logits are confidently wrong, and with a larger rate still the numbers overflow	Divide the learning rate by 10 and rerun
Validation far below training for many epochs	Something differs between the two passes: dropout or batch-norm in the wrong mode, or a validation set that is easier or leaked	Check model.eval() and the split

Task 7 produces the second row on purpose, and Task 8 the fourth and the third.

7. Make it overfit, on purpose

Five epochs on 50,000 images did not get to the turn. The quickest way there is not more epochs on the same data; it is less data and more steps on it. With 5,000 training images (40 steps per epoch instead of 391) and a learning rate of 0.5, sixty epochs are 2,400 steps, about six epochs’ worth of the full set, and the network can memorise the images it sees:

RunnableAll tracks

overfit: a tenth of the data, sixty epochs, learning rate 0.5

python train-mnist.py --epochs 60 --lr 0.5 --train-size 5000 --label overfit-5k --checkpoint mnist-overfit-5k.pt --labbook labbook.md

Output — what you should see

overfit-5k, author's CPU run, abridged

torch 2.14.0+cpu
device: cpu
data: train 5,000  val 10,000  test 10,000  batch 128 -> 40 steps per epoch
parameters: 203,530
  layers.0.weight      shape (256, 784)
  layers.0.bias        shape (256,)
  layers.2.weight      shape (10, 256)
  layers.2.bias        shape (10,)
epoch  0  (untrained)          val loss 2.3091  val acc 0.0645
epoch  1  train loss 1.0804  val loss 0.7966  val acc 0.7034    6.3 s  (saved)
epoch  2  train loss 0.4340  val loss 0.4270  val acc 0.8664    6.8 s  (saved)
epoch  3  train loss 0.3190  val loss 0.4005  val acc 0.8808    6.4 s  (saved)
epoch  4  train loss 0.2549  val loss 0.4769  val acc 0.8549    6.5 s
epoch  5  train loss 0.2253  val loss 0.3920  val acc 0.8790    6.7 s  (saved)
epoch  6  train loss 0.1854  val loss 0.2830  val acc 0.9172    6.5 s  (saved)
epoch  7  train loss 0.1572  val loss 0.3736  val acc 0.8880    6.6 s
epoch  8  train loss 0.1403  val loss 1.1236  val acc 0.7014    7.0 s
epoch  9  train loss 0.1926  val loss 2.6640  val acc 0.4916    6.5 s
epoch 10  train loss 0.3093  val loss 0.3507  val acc 0.8941    6.6 s
epoch 11  train loss 0.1162  val loss 0.5748  val acc 0.8568    6.4 s
epoch 12  train loss 0.1012  val loss 1.1779  val acc 0.7851    6.3 s
epoch 13  train loss 0.1268  val loss 0.2205  val acc 0.9370    6.2 s  (saved)
epoch 14  train loss 0.0672  val loss 0.2168  val acc 0.9391    7.0 s  (saved)
  ... epochs 15 to 23 omitted ...
epoch 24  train loss 0.0179  val loss 0.2080  val acc 0.9436    7.0 s  (saved)
epoch 25  train loss 0.0161  val loss 0.2114  val acc 0.9429    7.0 s
epoch 26  train loss 0.0144  val loss 0.2099  val acc 0.9425    6.1 s
  ... epochs 27 to 39 omitted ...
epoch 40  train loss 0.0065  val loss 0.2228  val acc 0.9430    6.8 s
  ... epochs 41 to 59 omitted ...
epoch 60  train loss 0.0030  val loss 0.2316  val acc 0.9439    7.2 s
best epoch 24: test loss 0.1911  test acc 0.9461  (404 s total, checkpoint mnist-overfit-5k.pt)
recorded in labbook.md

Validation loss reaches its minimum at epoch 24 and then creeps up, from 0.208 to 0.232 at epoch 60, while training loss keeps falling from 0.018 to 0.003, sixty times below where the first run stopped. (saved) stops appearing at the minimum, and the script tests that checkpoint rather than the last epoch, which is why its test accuracy is close to the first run’s despite a tenth of the data. Both curves are noisy between epochs 8 and 12 (validation loss spikes to 2.66 at epoch 9), which is what a learning rate near the edge of what plain SGD tolerates looks like before it settles. On a bigger model with more epochs the same divergence becomes the wide gap the previous lesson drew. Plot it:

RunnableAll tracks

plot the last two runs side by side

python plot-curves.py --labbook labbook.md --out curves.png --last 2

Record: the epoch of the last (saved) (the early-stopping point), the minimum validation loss, the validation loss at epoch 60, the final training loss, and the test accuracy.

If you have an accelerator and five more minutes, the same learning rate on the full set turns the same way, later: in the author’s run below the validation minimum came at epoch 13 and the kept checkpoint scored 0.981 on the test set. Ten times the data moved the turn from 960 steps (24 × 40) to 5,083 (13 × 391) and lifted test accuracy by three and a half points; that trade is the argument of the pretraining lesson in Part 3 at a scale of trillions of tokens.

RunnableAll tracks

optional: the full set, forty epochs, learning rate 0.5

python train-mnist.py --epochs 40 --lr 0.5 --label overfit-full --checkpoint mnist-overfit.pt --labbook labbook.md

Output — what you should see

overfit-full, author's CPU run (four threads), abridged

torch 2.14.0+cpu
device: cpu
data: train 50,000  val 10,000  test 10,000  batch 128 -> 391 steps per epoch
parameters: 203,530
  layers.0.weight      shape (256, 784)
  layers.0.bias        shape (256,)
  layers.2.weight      shape (10, 256)
  layers.2.bias        shape (10,)
epoch  0  (untrained)          val loss 2.3091  val acc 0.0645
epoch  1  train loss 0.3526  val loss 0.1804  val acc 0.9467   14.1 s  (saved)
  ... epochs 2 to 4 omitted ...
epoch  5  train loss 0.0604  val loss 0.0846  val acc 0.9752   12.0 s  (saved)
  ... epochs 6 to 11 omitted ...
epoch 12  train loss 0.0157  val loss 0.0714  val acc 0.9792   16.7 s  (saved)
epoch 13  train loss 0.0129  val loss 0.0688  val acc 0.9809    3.8 s  (saved)
epoch 14  train loss 0.0111  val loss 0.0698  val acc 0.9809    5.1 s
  ... epochs 15 to 19 omitted ...
epoch 20  train loss 0.0045  val loss 0.0734  val acc 0.9813   16.6 s
  ... epochs 21 to 39 omitted ...
epoch 40  train loss 0.0013  val loss 0.0795  val acc 0.9826    8.2 s
best epoch 13: test loss 0.0659  test acc 0.9810  (603 s total, checkpoint mnist-overfit.pt)
recorded in labbook.md

8. Break it, on purpose

Two short runs to see the failures the first lesson described, each with its own checkpoint name (the callout below says why):

RunnableAll tracks

learning rate 200 times too large

python train-mnist.py --epochs 3 --lr 20 --label lr-too-big --checkpoint mnist-lr20.pt --labbook labbook.md

Output — what you should see

lr-too-big, author's CPU run

torch 2.14.0+cpu
device: cpu
data: train 50,000  val 10,000  test 10,000  batch 128 -> 391 steps per epoch
parameters: 203,530
  layers.0.weight      shape (256, 784)
  layers.0.bias        shape (256,)
  layers.2.weight      shape (10, 256)
  layers.2.bias        shape (10,)
epoch  0  (untrained)          val loss 2.3091  val acc 0.0645
epoch  1  train loss 395.5690  val loss 2.8467  val acc 0.0990   37.4 s  (saved)
epoch  2  train loss 2.7856  val loss 2.8141  val acc 0.0956   69.8 s  (saved)
epoch  3  train loss 2.7217  val loss 2.9581  val acc 0.1057   76.7 s
best epoch 2: test loss 2.7254  test acc 0.0978  (184 s total, checkpoint mnist-lr20.pt)
recorded in labbook.md

RunnableAll tracks

learning rate 1,000 times too small

python train-mnist.py --epochs 3 --lr 0.0001 --label lr-too-small --checkpoint mnist-lr-tiny.pt --labbook labbook.md

Output — what you should see

lr-too-small, author's CPU run

torch 2.14.0+cpu
device: cpu
data: train 50,000  val 10,000  test 10,000  batch 128 -> 391 steps per epoch
parameters: 203,530
  layers.0.weight      shape (256, 784)
  layers.0.bias        shape (256,)
  layers.2.weight      shape (10, 256)
  layers.2.bias        shape (10,)
epoch  0  (untrained)          val loss 2.3091  val acc 0.0645
epoch  1  train loss 2.3053  val loss 2.3015  val acc 0.0789   73.6 s  (saved)
epoch  2  train loss 2.2976  val loss 2.2938  val acc 0.0997   70.9 s  (saved)
epoch  3  train loss 2.2900  val loss 2.2862  val acc 0.1294   77.8 s  (saved)
best epoch 3: test loss 2.2856  test acc 0.1352  (222 s total, checkpoint mnist-lr-tiny.pt)
recorded in labbook.md

The first run’s training loss in epoch 1 is in the hundreds: each step overshoots so far that the logits are enormous and confidently wrong, and -log of a near-zero probability is a large number. From epoch 2 it sits between 2.7 and 3.0, above the 2.30 of guessing evenly, because a network that is sure and wrong pays more than one that shrugs; the kept checkpoint scores 0.0978 on the test set, a ten-sided coin. A larger rate still overflows to nan, which never counts as an improvement, so the script saves nothing and says so. The second run is the opposite failure: everything moves the right way, by 0.008 per epoch, so after three epochs it is still at 2.29 where the default run passed 0.61 after one. Neither is a bug; both are the learning rate. Record for each: the label, the epoch-1 and final training loss, the test accuracy, and whether anything was saved. Then plot every run, so that the four shapes sit side by side:

RunnableAll tracks

plot every run: healthy, overfit, diverged, starved

python plot-curves.py --labbook labbook.md --out curves.png

Output — what you should see

plot-curves.py after four runs, author's run

run  label                  lr epochs  train best min val loss final train
1    first-run             0.1      5  50000    5       0.1937      0.1883
2    overfit-5k            0.5     60   5000   24       0.2080      0.0030
3    lr-too-big           20.0      3  50000    2       2.8141      2.7217
4    lr-too-small       0.0001      3  50000    3       2.2862      2.2900
wrote curves.png (4 run(s) from labbook.md)

With the optional overfit-full run there is a fifth row, 5 overfit-full 0.5 40 50000 13 0.0688 0.0013, and the last line says 5 run(s). The four panels are, in order, rows 1, 2, 4 and 3 of the shape table in Task 6. Open curves.png, or rerun cell 2 of the notebook.

✕A checkpoint that is not the run you think

If every run wrote the default mnist_mlp.pt, the slow run in this task would overwrite the good checkpoint from Task 5 with three epochs of barely trained weights, and a later reload-mnist.py would evaluate weights at 0.13 test accuracy while the notebook says 0.9465. The symptom in a real project is a model that “got worse after I retrained”: the file changed under a name that did not. Name checkpoints by run, as every command on this page does with --checkpoint, and record the name in the notebook next to the number.

9. Save and reload the weights

The first run saved its best weights to mnist-first.pt with torch.save(model.state_dict(), ...). A state_dict is, in the PyTorch tutorial’s words, “a Python dictionary object that maps each layer to its parameter tensor”; the file is a zip archive (the format torch.save has used since PyTorch 1.6) holding one entry per tensor’s bytes and a small pickle that names them. The script below opens the file both ways, rebuilds the network from its class, loads the weights, evaluates the test set once with no training, and finally loads the same file into a network of the wrong width to show what a mismatch looks like.

RunnableAll tracks

reload-mnist.py

"""Reload a checkpoint written by train-mnist.py, inspect it, and evaluate it without training.

Purpose: prove that the .pt file is the model. List what the file holds and how many bytes
         each tensor is, rebuild the network from its class, load the weights, evaluate the
         test set once, and then show what a shape mismatch looks like.
Platform: all (runs on the CPU on purpose: a checkpoint is device-independent)
Minimum memory: 8 GB
Assumes: torch and torchvision are installed; a checkpoint from train-mnist.py exists;
         MNIST is already under ./data (train-mnist.py downloaded it).

Usage: python reload-mnist.py [--checkpoint mnist_mlp.pt] [--hidden 256] [--data ./data]
"""
import argparse
import sys
import zipfile
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class TwoLayerNet(nn.Module):
    """The same architecture as train-mnist.py: the state_dict keys must match it exactly."""

    def __init__(self, hidden: int = 256) -> None:
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(nn.Linear(28 * 28, hidden), nn.ReLU(), nn.Linear(hidden, 10))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(self.flatten(x))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--checkpoint", default="mnist_mlp.pt")
    parser.add_argument("--hidden", type=int, default=256, help="must match the run that wrote the checkpoint")
    parser.add_argument("--data", default="./data")
    args = parser.parse_args()

    path = Path(args.checkpoint)
    if not path.exists():
        sys.exit(f"reload-mnist: {path} does not exist; run train-mnist.py first")

    # 1. The file is a zip archive (torch.save's format since PyTorch 1.6): names and sizes.
    print(f"file: {path}  {path.stat().st_size:,} bytes on disk")
    with zipfile.ZipFile(path) as zf:
        for info in zf.infolist():
            print(f"  {info.filename:40s} {info.file_size:>9,} bytes")

    # 2. The state_dict: an ordered mapping from parameter name to tensor.
    sd = torch.load(path, map_location="cpu")  # weights_only=True is the default: tensors only, no code
    print(f"\nstate_dict: {type(sd).__name__} with {len(sd)} entries")
    total = 0
    for name, t in sd.items():
        nbytes = t.numel() * t.element_size()
        total += nbytes
        print(f"  {name:20s} {str(tuple(t.shape)):12s} {str(t.dtype):14s} {t.numel():>8,} x {t.element_size()} B = {nbytes:>9,} bytes")
    print(f"  {'total':20s} {'':12s} {'':14s} {sum(t.numel() for t in sd.values()):>8,} params   {total:>9,} bytes")

    # 3. Rebuild the network and load the weights into it.
    model = TwoLayerNet(args.hidden)
    try:
        result = model.load_state_dict(sd)  # strict=True: every key and shape must match
    except RuntimeError as err:
        sys.exit("reload-mnist: the checkpoint was written with a different --hidden; "
                 + str(err).splitlines()[1].strip())
    model.eval()
    print(f"\nload_state_dict: missing {list(result.missing_keys)}  unexpected {list(result.unexpected_keys)}")

    # 4. Evaluate the test set once, exactly as train-mnist.py did, without a single training step.
    test_set = datasets.MNIST(args.data, train=False, download=False, transform=transforms.ToTensor())
    loader = DataLoader(test_set, batch_size=1000)
    loss_fn = nn.CrossEntropyLoss()
    total_loss, correct, seen = 0.0, 0, 0
    with torch.no_grad():
        for images, labels in loader:
            logits = model(images)
            total_loss += loss_fn(logits, labels).item() * labels.size(0)
            correct += (logits.argmax(dim=1) == labels).sum().item()
            seen += labels.size(0)
    print(f"test loss {total_loss / seen:.4f}  test acc {correct / seen:.4f}  ({seen:,} images, cpu, no training)")

    # 5. What a mismatch looks like: the same file into a network of a different width.
    wrong_hidden = 128 if args.hidden != 128 else 64
    print(f"\nloading the same file into TwoLayerNet(hidden={wrong_hidden}) ...")
    try:
        TwoLayerNet(wrong_hidden).load_state_dict(sd)
    except RuntimeError as err:
        first_lines = str(err).splitlines()[:2]
        print("RuntimeError:", first_lines[0])
        if len(first_lines) > 1:
            print("   ", first_lines[1].strip()[:110], "...")


if __name__ == "__main__":
    main()

Download reload-mnist.py99 lines

RunnableAll tracks

reload and evaluate without training

python reload-mnist.py --checkpoint mnist-first.pt

Output — what you should see

reload-mnist.py, author's run

file: mnist-first.pt  816,709 bytes on disk
  mnist-first/data.pkl                           576 bytes
  mnist-first/.format_version                      1 bytes
  mnist-first/.storage_alignment                   2 bytes
  mnist-first/byteorder                            6 bytes
  mnist-first/data/0                         802,816 bytes
  mnist-first/data/1                           1,024 bytes
  mnist-first/data/2                          10,240 bytes
  mnist-first/data/3                              40 bytes
  mnist-first/version                              2 bytes
  mnist-first/.data/serialization_id              40 bytes
state_dict: OrderedDict with 4 entries
  layers.0.weight      (256, 784)   torch.float32   200,704 x 4 B =   802,816 bytes
  layers.0.bias        (256,)       torch.float32       256 x 4 B =     1,024 bytes
  layers.2.weight      (10, 256)    torch.float32     2,560 x 4 B =    10,240 bytes
  layers.2.bias        (10,)        torch.float32        10 x 4 B =        40 bytes
  total                                             203,530 params     814,120 bytes
load_state_dict: missing []  unexpected []
test loss 0.1790  test acc 0.9465  (10,000 images, cpu, no training)
loading the same file into TwoLayerNet(hidden=128) ...
RuntimeError: Error(s) in loading state_dict for TwoLayerNet:
    size mismatch for layers.0.weight: copying a param with shape torch.Size([256, 784]) from checkpoint, the shap ...

Four tensors, in FP32, four bytes per number: 203,530 × 4 = 814,120 bytes of weights, in a file of 816,709 bytes, the extra 2,589 being the zip directory, the entry names (which include the file’s own name, so a longer checkpoint name makes a slightly larger file) and the pickle that names the tensors. This is the “bytes per parameter times parameters” arithmetic of the precision lesson on a model small enough to check by hand, and the same arithmetic sizes every checkpoint in the course:

Model	Parameters	Bytes per parameter	Weights on disk (arithmetic, not a measurement)
This network, FP32	203,530	4	814,120 bytes
This network, BF16	203,530	2	407,060 bytes
Same with --hidden 512, FP32	407,050	4	1,628,200 bytes
An 8B-parameter model, BF16	8,000,000,000	2	16 GB
An 8B-parameter model, 4-bit	8,000,000,000	0.5	4 GB, plus the scales the format adds

The test accuracy equals the training run’s to the fourth decimal: the file is the model, and nothing about it lives outside the file except its class definition, which is why the keys layers.0.weight and so on must match the module names exactly. The last lines are the failure mode: load_state_dict is strict by default and refuses a tensor of the wrong shape with a size mismatch error naming the key and both shapes, which is what loading a checkpoint into the wrong architecture looks like in Part 13.

⛨Security

torch.load unpickles. Its documentation says: “Never load data from an untrusted source.” The default weights_only=True restricts the unpickler to tensors and primitive types, which is why the script does not pass anything else, but a .pt file from someone you do not trust is still code you are running. Models from Part 2 onwards arrive as .safetensors, a format with no code in it, for exactly this reason.

Add the third cell to the notebook, which looks at what the first layer learned: each of the 256 rows of layers.0.weight is a 784-vector, and reshaped to 28×28 it is the pattern of pixels that hidden unit responds to:

Fragment — not complete on its own

part-01.ipynb, cell 3

import torch, matplotlib.pyplot as plt
sd = torch.load("mnist-first.pt", map_location="cpu")
w = sd["layers.0.weight"]                      # shape (256, 784)
fig, axes = plt.subplots(2, 8, figsize=(12, 3))
for i, ax in enumerate(axes.flat):
    ax.imshow(w[i].reshape(28, 28), cmap="gray")
    ax.set_title(f"unit {i}", fontsize=8); ax.axis("off")
plt.show()

Sixteen blurry blobs and strokes, some looking like partial digits: that is what 391 × 5 steps of gradient descent made of random noise. Save the notebook. What to record: the byte total (814,120) and the file size, next to the first run’s line.

10. Track M only: the same model in MLX

MLX is Apple’s array framework and the native way to train and run models on Apple silicon; Parts 8 and 13 use it heavily. The script below is the PyTorch one rewritten in MLX, with the same options and output format, so that the differences are visible line by line:

Step	PyTorch (train-mnist.py)	MLX (train-mnist-mlx.py)
Where tensors live	images.to(device) moves each batch to the GPU	Nowhere to move to: arrays sit in unified memory and the default device runs the operations
Parameters exist	On construction	Lazily; mx.eval(model.parameters()) materialises them
Gradients	loss.backward() writes .grad on each parameter	nn.value_and_grad(model, loss_fn) returns a function giving (loss, grads) as a tree
Update	optimiser.step()	optimiser.update(model, grads) then mx.eval(model.parameters(), optimiser.state) to force the lazy graph to run
Evaluation mode	model.eval() and torch.no_grad()	Nothing to switch: no dropout here, and no gradient is computed unless asked for
Checkpoint	torch.save(state_dict) to a .pt zip	model.save_weights("....safetensors"); .npz also accepted
Choosing a device	--device cpu	--device cpu calls mx.set_default_device(mx.cpu)

RunnableTrack M · Apple silicon

train-mnist-mlx.py

"""Train the same two-layer MNIST network with MLX, Apple's array framework.

Purpose: the Track M native path for the first training run, kept line for line as close
         to train-mnist.py as the two frameworks allow, so the differences are visible:
         arrays live in unified memory, computation is lazy until mx.eval, and gradients
         come from a function transform (nn.value_and_grad) rather than from .backward().
Platform: mac (Apple silicon, macOS 14 or later, native arm64 Python, mlx installed);
          also runs on the CPU build of MLX on Linux (pip install "mlx[cpu]") for checking.
Minimum memory: 8 GB
Assumes: mlx and torchvision are installed in the active environment (torchvision is used
         only to download and decode MNIST), and about 70 MB of free disk under ./data.

Usage: python train-mnist-mlx.py [--epochs 5] [--lr 0.1] [--batch-size 128] [--hidden 256]
                                 [--train-size 50000] [--device gpu|cpu]
                                 [--checkpoint mnist_mlp.safetensors] [--label mlx-first-run]
                                 [--labbook labbook.md] [--seed 0] [--data ./data]

The output format is the same as train-mnist.py so the two logs can be compared column
by column; the checkpoint is a .safetensors file, which MLX writes natively.
"""
import argparse
import json
import math
import sys
import time
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
from mlx.utils import tree_flatten
from torchvision import datasets


def load_split(root: str, train: bool) -> tuple[mx.array, mx.array]:
    ds = datasets.MNIST(root, train=train, download=True)
    images = ds.data.numpy().astype("float32") / 255.0      # bytes 0..255 -> 0.0..1.0
    labels = ds.targets.numpy().astype("int32")
    return mx.array(images.reshape(len(images), -1)), mx.array(labels)   # (n, 784), (n,)


class TwoLayerNet(nn.Module):
    """784 inputs -> hidden units with ReLU -> 10 outputs, one per digit."""

    def __init__(self, hidden: int = 256) -> None:
        super().__init__()
        self.l1 = nn.Linear(28 * 28, hidden)
        self.l2 = nn.Linear(hidden, 10)

    def __call__(self, x: mx.array) -> mx.array:
        return self.l2(nn.relu(self.l1(x)))


def loss_fn(model: TwoLayerNet, x: mx.array, y: mx.array) -> mx.array:
    return mx.mean(nn.losses.cross_entropy(model(x), y))


def evaluate(model: TwoLayerNet, x: mx.array, y: mx.array) -> tuple[float, float]:
    logits = model(x)
    loss = mx.mean(nn.losses.cross_entropy(logits, y)).item()
    acc = mx.mean(mx.argmax(logits, axis=1) == y).item()
    return loss, acc


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--lr", type=float, default=0.1, help="learning rate for plain SGD")
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--hidden", type=int, default=256)
    parser.add_argument("--train-size", type=int, default=50_000,
                        help="how many of the 50,000 training images to train on")
    parser.add_argument("--device", default="gpu", choices=["gpu", "cpu"],
                        help="MLX default device; gpu is the Apple GPU through Metal")
    parser.add_argument("--checkpoint", default="mnist_mlp.safetensors")
    parser.add_argument("--label", default="", help="a short name for this run in the notebook")
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--data", default="./data")
    args = parser.parse_args()
    if not 1 <= args.train_size <= 50_000:
        sys.exit("train-mnist-mlx: --train-size must be between 1 and 50000")
    if not args.checkpoint.endswith((".safetensors", ".npz")):
        sys.exit("train-mnist-mlx: --checkpoint must end in .safetensors or .npz (the formats save_weights writes)")

    if args.device == "cpu":
        mx.set_default_device(mx.cpu)
    elif mx.default_device() != mx.gpu:
        sys.exit(f"train-mnist-mlx: the default device is {mx.default_device()}, not the GPU; "
                 "pass --device cpu on a build without Metal")
    mx.random.seed(args.seed)
    np.random.seed(args.seed)
    print(f"mlx {mx.__version__}")
    print(f"device: {mx.default_device()}")

    x_all, y_all = load_split(args.data, train=True)
    x_test, y_test = load_split(args.data, train=False)
    x_train, y_train = x_all[:args.train_size], y_all[:args.train_size]   # the first 50,000 (or fewer)
    x_val, y_val = x_all[50_000:], y_all[50_000:]                          # the last 10,000 held out
    n = x_train.shape[0]
    steps = math.ceil(n / args.batch_size)
    print(f"data: train {n:,}  val {x_val.shape[0]:,}  test {x_test.shape[0]:,}  "
          f"batch {args.batch_size} -> {steps} steps per epoch")

    model = TwoLayerNet(args.hidden)
    mx.eval(model.parameters())   # parameters are lazy until evaluated
    n_params = sum(p.size for _, p in tree_flatten(model.parameters()))
    print(f"parameters: {n_params:,}")
    for name, p in tree_flatten(model.parameters()):
        print(f"  {name:20s} shape {tuple(p.shape)}")

    optimiser = optim.SGD(learning_rate=args.lr)
    step = nn.value_and_grad(model, loss_fn)   # a function returning (loss, grads) for the model's parameters

    val_loss, val_acc = evaluate(model, x_val, y_val)
    print(f"epoch  0  (untrained)          val loss {val_loss:.4f}  val acc {val_acc:.4f}")

    history, best_val, best_epoch = [], float("inf"), 0
    started = time.time()
    for epoch in range(1, args.epochs + 1):
        epoch_started = time.time()
        order = np.random.permutation(n)
        running = 0.0
        for start in range(0, n, args.batch_size):
            idx = mx.array(order[start:start + args.batch_size])
            loss, grads = step(model, x_train[idx], y_train[idx])   # forward + loss + backward
            optimiser.update(model, grads)                           # step
            mx.eval(model.parameters(), optimiser.state)             # force the lazy graph to run
            running += loss.item() * idx.shape[0]
        train_loss = running / n
        val_loss, val_acc = evaluate(model, x_val, y_val)
        seconds = time.time() - epoch_started
        history.append({"epoch": epoch, "train_loss": round(train_loss, 4), "val_loss": round(val_loss, 4),
                        "val_acc": round(val_acc, 4), "seconds": round(seconds, 1)})
        marker = ""
        if val_loss < best_val:
            best_val, best_epoch = val_loss, epoch
            model.save_weights(args.checkpoint)
            marker = "  (saved)"
        print(f"epoch {epoch:2d}  train loss {train_loss:.4f}  val loss {val_loss:.4f}  "
              f"val acc {val_acc:.4f}  {seconds:5.1f} s{marker}")
    elapsed = time.time() - started

    if best_epoch == 0:
        test_loss, test_acc = float("nan"), float("nan")
        print(f"no epoch improved the validation loss: nothing saved to {args.checkpoint}, "
              f"test set not used  ({elapsed:.0f} s)")
    else:
        model.load_weights(args.checkpoint)
        test_loss, test_acc = evaluate(model, x_test, y_test)
        print(f"best epoch {best_epoch}: test loss {test_loss:.4f}  test acc {test_acc:.4f}  "
              f"({elapsed:.0f} s total, checkpoint {args.checkpoint})")

    if args.labbook:
        record = {
            "lab": "part-01/train-mnist-mlx", "label": args.label, "device": str(mx.default_device()),
            "mlx": mx.__version__, "epochs": args.epochs, "lr": args.lr, "batch_size": args.batch_size,
            "hidden": args.hidden, "train_size": n, "parameters": n_params,
            "best_epoch": best_epoch, "test_acc": None if math.isnan(test_acc) else round(test_acc, 4),
            "seconds": round(elapsed, 1), "checkpoint": args.checkpoint, "history": history,
        }
        with Path(args.labbook).open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record) + "\n")
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()

Download train-mnist-mlx.py169 lines

RunnableTrack M · Apple silicon

MLX run

python train-mnist-mlx.py --epochs 5 --label mlx-first-run --labbook labbook.md

The output below is from the Linux CPU build of MLX (pip install "mlx[cpu]", which the MLX install page documents) on the author’s machine, used to check the script; on a Mac the device line reads Device(gpu, 0) and the seconds column is the Apple GPU’s:

Output — what you should see

train-mnist-mlx.py on the Linux CPU build of MLX, author's run

mlx 0.32.2
device: Device(cpu, 0)
data: train 50,000  val 10,000  test 10,000  batch 128 -> 391 steps per epoch
parameters: 203,530
  l1.weight            shape (256, 784)
  l1.bias              shape (256,)
  l2.weight            shape (10, 256)
  l2.bias              shape (10,)
epoch  0  (untrained)          val loss 2.3076  val acc 0.1004
epoch  1  train loss 0.6144  val loss 0.3149  val acc 0.9126   14.3 s  (saved)
epoch  2  train loss 0.3058  val loss 0.2619  val acc 0.9248   16.4 s  (saved)
epoch  3  train loss 0.2529  val loss 0.2169  val acc 0.9412   14.2 s  (saved)
epoch  4  train loss 0.2166  val loss 0.1929  val acc 0.9475   13.5 s  (saved)
epoch  5  train loss 0.1893  val loss 0.1744  val acc 0.9526   14.7 s  (saved)
best epoch 5: test loss 0.1779  test acc 0.9466  (73 s total, checkpoint mnist_mlp.safetensors)
recorded in labbook.md

The same shapes, loss values within run-to-run noise of the PyTorch run (the initialisation and the shuffling differ, so the digits do not match exactly), and similar accuracy. Then the same job on the CPU, your first measurement of one job on two devices:

RunnableTrack M · Apple silicon

MLX on the CPU, and PyTorch on the CPU, for comparison

python train-mnist-mlx.py --epochs 5 --device cpu --label mlx-cpu --checkpoint mnist_mlp_cpu.safetensors --labbook labbook.md
python train-mnist.py --epochs 5 --device cpu --label torch-cpu --checkpoint mnist-torch-cpu.pt --labbook labbook.md

Record all four seconds-per-epoch figures in one table in the notebook: PyTorch MPS, PyTorch CPU, MLX GPU, MLX CPU. Do not expect the GPU to win by much on a 203,530-parameter model: the work per step is tiny and the overhead of launching it dominates, the first instance of a lesson Part 5 makes with numbers: the accelerator pays off when the arithmetic per byte moved is large.

Owners of the other tracks can do the same two-device comparison with PyTorch alone, --device cpu against the default, and record the two seconds-per-epoch figures.

Checkpoint review before you mark the environment usable

Use the following sequence to distinguish installation success from a working learning loop. Run only the track you selected; the other tabs are alternatives, not additional installation steps.

Checkpoint	Evidence to retain	Stop and investigate when
Python environment	Python executable, package versions and selected device	The executable belongs to a different environment or the device is an unintended CPU fallback
Baseline training	Complete training log and both loss curves	Loss becomes non-finite, the training set is empty or the target labels are wrong
Deliberate overfit	Baseline and reduced-data curves under distinct run names	The two experiments changed several settings and cannot isolate the intended effect
Saved model	Checkpoint path and reload output on the held-out images	Reload creates a new random model or uses different preprocessing
Interpretation	One explanation of the train/validation gap	The conclusion relies only on training accuracy

After reloading, use the same normalisation and label mapping as training. A successfully read file is not enough: compare predictions or held-out metrics from the saved weights. Keep the original baseline before deliberately breaking the run, and give every altered experiment its own output name. The intentionally failed run is complete when you can connect its observed symptom to the change you made, restore the baseline and explain the difference.

Validation

You are done when every line below passes. Run them from ~/llm-course with the environment active (or in the container shell on Track S).

RunnableAll tracks

validation

python -c "import torch; print(torch.__version__)"
python - <<'PY'
import json
runs = [json.loads(l) for l in open("labbook.md") if l.startswith("{")]
print(len(runs), "runs recorded:", [r.get("label") for r in runs])
PY
ls -l mnist-first.pt mnist-overfit-5k.pt curves.png part-01.ipynb
python reload-mnist.py --checkpoint mnist-first.pt | grep -E "total|test acc"

Output — what you should see

2.14.0+cu130
4 runs recorded: ['first-run', 'overfit-5k', 'lr-too-big', 'lr-too-small']
# 5 with overfit-full; more on Track M
-rw-r--r-- 1 you you 816709 Sep 12 12:xx mnist-first.pt
-rw-r--r-- 1 you you 816759 Sep 12 12:xx mnist-overfit-5k.pt
-rw-r--r-- 1 you you xxxxxx Sep 12 12:xx curves.png
-rw-r--r-- 1 you you xxxxx  Sep 12 12:xx part-01.ipynb
  total                                       203,530 params     814,120 bytes
test loss 0.xxxx  test acc 0.94xx  (10,000 images, cpu, no training)

Check	Pass means
torch.__version__	Carries the suffix your track expects from the table in Task 3 (+cu130, +rocm7.2, plain 2.14.0 on a Mac), or +cpu with the reason in the notebook
Runs recorded	At least four labels: first-run, overfit-5k, lr-too-big, lr-too-small; five with overfit-full; more on Track M
labbook.md Machine and Environment	Every field filled; the accelerator: line copied verbatim
first-run	Test accuracy 0.94 or better; best_epoch 5
overfit-5k	best_epoch well before 60; validation loss in history rises after it while training loss keeps falling to below 0.01
lr-too-big	Training loss far above 2.30 in epoch 1 and not below it afterwards, or nan; test accuracy near 0.10, or nothing saved
lr-too-small	Training loss after three epochs still above 2.2, that is, above the first run’s epoch-1 value of 0.61
reload-mnist.py	814,120 bytes, and a test accuracy equal to the first run’s
curves.png and part-01.ipynb	Exist in the course directory; curves.png has four panels (five with the optional run); the notebook has three cells that ran

Expected outcome

Done means every Validation check passes and ~/llm-course holds .venv, labbook.md, curves.png, part-01.ipynb and the checkpoints mnist-first.pt, mnist-overfit-5k.pt, mnist-lr20.pt and mnist-lr-tiny.pt.

Troubleshooting

Symptoms are quoted as the tools print them, where they were captured on 2026-09-12; wording may differ by version where noted.

Symptom	Cause	Fix
uv: command not found right after installing	The installer edited your shell profile, but this shell was started before that	Open a new terminal, or source ~/.bashrc (~/.zshrc on a Mac)
error: Failed to create virtual environment / Caused by: A virtual environment already exists at: .venv / hint: Use the --clear flag or set UV_VENV_CLEAR=1 to replace the existing virtual environment	.venv exists from an earlier attempt	FRESH=1 TRACK=... bash setup-env.sh (which passes --clear), or just rerun without FRESH to reuse it
ModuleNotFoundError: No module named 'torch'	The environment is not active, or python is another interpreter	source .venv/bin/activate, then which python must print .../llm-course/.venv/bin/python
ModuleNotFoundError: No module named 'torchvision' or 'matplotlib' in the NGC container	The image tag you pulled does not ship it (the 26.08 release notes list neither)	pip install torchvision matplotlib inside the container; note the versions in the Environment section
accelerator: none visible on Track N, and nvidia-smi fails	No driver, or the wrong one: on Linux the distribution’s NVIDIA driver package is missing; on Windows the driver is not installed, or a Linux driver was installed inside WSL against the guide’s rule	Linux: install the driver package and reboot. Windows: install the NVIDIA Windows driver, wsl.exe --update, and run /usr/lib/wsl/lib/nvidia-smi inside WSL
accelerator: none visible on Track N, nvidia-smi works, torch 2.14.0+cu130	The driver is older than 580, the CUDA release notes’ minimum for CUDA 13.x	TORCH_BACKEND=cu126 FRESH=1 TRACK=nvidia bash setup-env.sh (driver 560.28.03 or later), or update the driver
accelerator: none visible on Track N, torch 2.14.0+cpu	--torch-backend=auto found no CUDA driver to match and fell back to the CPU build	TORCH_BACKEND=cu130 FRESH=1 TRACK=nvidia bash setup-env.sh with a driver of 580 or later, TORCH_BACKEND=cu126 FRESH=1 TRACK=nvidia bash setup-env.sh with an older one
torch 2.11.0+cu128 printed by the verification	--torch-backend=auto chose the cu128 index for a 570-series driver, and that index stops at torch 2.11.0 (read 2026-09-12)	TORCH_BACKEND=cu126 FRESH=1 TRACK=nvidia bash setup-env.sh (driver 560.28.03 or later), or update the driver to 580 and rerun with TORCH_BACKEND=cu130
AssertionError: Torch not compiled with CUDA enabled	A CPU wheel was installed and something asked for cuda explicitly	Reinstall with the right backend, as above; train-mnist.py never asks for a device it cannot see unless --device says so
accelerator: none visible on Track X, rocminfo shows gfx1151	One of: user not in render and video (log out and in after usermod); wheel ROCm version does not match the installed ROCm (compare hip build: with amd-smi version); the stable wheel lacks this GPU	Fix groups; match versions; try ROCM_NIGHTLY=1; otherwise finish on the CPU, write “cpu, ROCm x.y wheel did not see gfx1151” in the notebook, and revisit in Part 5
rocminfo shows no gfx line at all	The ROCm runtime does not enumerate the GPU; AMD’s prerequisites page states that ROCm does not currently support integrated graphics	This is below PyTorch. Finish on the CPU; Part 5’s hardware page carries the ROCm status for this chip
RuntimeError: PyTorch is not linked with support for mps devices	The Python is an x86 build under Rosetta, or a Linux wheel	python -c "import platform; print(platform.machine())" must print arm64; if not, FRESH=1 TRACK=mac bash setup-env.sh from a native terminal
accelerator: none visible on Track M with arm64	macOS below 14	The MPS notes’ own check prints which case it is: torch.backends.mps.is_built() False means the wheel, True means the OS or device. Update macOS
train loss nan at the default learning rate, on an accelerator	The same run is healthy on the CPU (below), so the accelerator build mishandles an operation for this model	Rerun with --device cpu; if the CPU run is healthy, record it as a measurement-differs issue with your versions and finish on the CPU
MNIST download fails (HTTPError, URLError, or hangs)	A proxy or an offline machine; torchvision fetches from a mirror	Download on another machine and copy data/MNIST/raw/ across; the script needs the four files listed in the MNIST docs’ root description
jupyter lab says the port is in use	Another server, or a previous one still running	jupyter lab --port 8889, and stop the old one with Control-C in its terminal
The browser opens an empty page or asks for a token	You opened localhost:8888 without the token	Paste the full http://localhost:8888/lab?token=... line from the terminal
Track S container prints accelerator: none visible	The container was started without --gpus all, or the toolkit is not configured	nvidia-smi on the host first; then rerun the script, which passes --gpus all; the DGX Spark documentation’s own check is docker run -it --gpus=all nvcr.io/nvidia/cuda:13.0.1-devel-ubuntu24.04 nvidia-smi
permission denied while trying to connect to the Docker daemon socket	Your user is not in the docker group	sudo bash setup-env-spark.sh, or add yourself to the group and log in again, as NVIDIA’s page describes
RuntimeError: Error(s) in loading state_dict ... size mismatch for layers.0.weight	The checkpoint was written with a different --hidden	Pass the same --hidden to reload-mnist.py as to the training run; the error names both shapes

Cleanup

Leave the environment, the notebook file, the scripts and labbook.md; every later lab uses them.

⚠Warning

The block below deletes this lab’s checkpoints, curves.png and the downloaded MNIST files in ~/llm-course/data; rerunning the tasks recreates them. Track S: docker rmi nvcr.io/nvidia/pytorch:26.08-py3 removes the 11 GB image, which later Spark labs pull again, so leave it unless the disk is needed now. Nothing on this page removes anything you did not create in this lab.

RunnableAll tracks

optional: reclaim the disk this lab used

cd ~/llm-course
rm -f mnist-*.pt mnist_mlp*.safetensors curves.png
rm -rf data

What you learned

Objective	The observation that proved it	Recorded
The environment is reproducible and isolated	pyvenv.cfg names the interpreter; which python points into .venv; include-system-site-packages = false	Python, torch, mlx, jupyterlab, matplotlib versions in Environment
The accelerator is a wheel you chose, then a device you check	torch.__version__ carries the build suffix; the accelerator: line names the device; on X the name is cuda with hip build set	The accelerator: line verbatim; the index or TORCH_BACKEND used
Scripts record, notebooks look	Every run is a JSON line with its settings; the notebook’s three cells only read what the scripts wrote	part-01.ipynb exists; four or more labelled JSON lines
The training loop is real	Forward, loss, backward, step in four marked lines; loss from 2.30 (ln 10) to 0.19 in five epochs	first-run: best epoch, test accuracy, seconds per epoch
A loss curve and a validation curve are read, not admired	The four-run curves.png: healthy, overfit (validation turns at epoch 24 while training falls to 0.003), diverged (hundreds, then above 2.30), starved (2.31 to 2.29 in three epochs)	overfit-5k: the turn epoch, minimum validation loss, final training loss
Weights are bytes	Four tensors × FP32 = 814,120 bytes inside an 816,709-byte zip; a fresh network loaded from it scores the same to four decimals; the wrong width raises size mismatch	The byte total; reload-mnist.py test accuracy

Where it goes next: Part 2’s lab reuses this environment and notebook to do the bytes arithmetic on a real model’s .safetensors files at 1.7B parameters; Part 5’s preparation lab records this machine’s measured bandwidth beside the versions you wrote down today, and its Strix Halo page carries the ROCm status for Track X; Part 11 builds the fine-tuning environment on the same uv habits; and the two curves come back in every training lab from Part 12 onwards, where “the validation curve turned at step N” decides which checkpoint ships. Where it came from: the loop, the loss and the learning rate are What Learning Means; the two-layer network is Neural Networks, Activations and Backpropagation; the split and the two curves are Generalisation; and bytes per parameter is Tensors, GPUs and Precision.

Check your understanding

Question 1. On Track X the verification printed "accelerator: cuda / AMD Radeon Graphics" with "hip build: 7.2.x". Is something wrong?

• Yes, the NVIDIA build was installed by mistake
• No: the ROCm build of PyTorch exposes AMD GPUs through the same "cuda" device name, and the hip build line is what identifies it
• Yes, ROCm must be reinstalled
• No, but the GPU will not be used

Show the answer and why

Answer: No: the ROCm build of PyTorch exposes AMD GPUs through the same "cuda" device name, and the hip build line is what identifies it

The ROCm build reports AMD GPUs as the cuda device; torch.version.hip carrying a version, and torch.version.cuda being None, is how you tell a ROCm wheel from a CUDA one. The GPU is in use.

Question 2. During the sixty-epoch run on 5,000 images, "(saved)" stopped appearing after epoch 24 while training loss kept falling to 0.003. Which checkpoint does the script test, and why?

• Epoch 60, because it has the lowest training loss
• Epoch 24, because validation loss was lowest there and the later epochs were fitting the training set's noise
• Epoch 1, because it is the least overfitted
• The average of all epochs

Show the answer and why

Answer: Epoch 24, because validation loss was lowest there and the later epochs were fitting the training set's noise

The script saves whenever validation loss improves and reloads that file before testing. That is early stopping, and it is why the test set is looked at only once, with the checkpoint the validation set chose rather than the one the training set liked.

Question 3. You rerun the first run with --hidden 512. How many parameters, and how many bytes will the FP32 checkpoint's tensors hold?

• 407,050 parameters and 1,628,200 bytes
• 407,050 parameters and 814,120 bytes, because the file size does not depend on width
• 203,530 parameters and 814,120 bytes, because MNIST is the same
• 406,528 parameters and 1,626,112 bytes

Show the answer and why

Answer: 407,050 parameters and 1,628,200 bytes

784 × 512 + 512 + 512 × 10 + 10 = 407,050, times four bytes. The last option forgets the two bias vectors. The same multiplication, with two bytes for BF16 or half a byte for four-bit formats, sizes every model in the course.

Question 4. Which of these three lines is the bug in a fresh course directory on Track N? (a) uv venv --python 3.12 .venv (b) source .venv/bin/activate (c) pip install torch torchvision

• (a): uv venv needs --seed to work at all
• (b): activation is not needed with uv
• (c): a uv environment has no pip unless --seed was used, and even where a pip exists it may not be the environment's; the course installs with uv pip install
• None; all three are fine

Show the answer and why

Answer: (c): a uv environment has no pip unless --seed was used, and even where a pip exists it may not be the environment's; the course installs with uv pip install

uv venv creates no pip by default, so (c) either fails with "pip: command not found" or, worse, runs a system pip and installs torch somewhere the environment cannot see. uv pip install always targets the active environment (or .venv in the current directory).

Question 5. The untrained network printed val loss 2.30. The same script adapted to a 100-class dataset prints, before training, a loss near:

• 2.30, because the network has not learned anything either way
• 4.61, because -ln(1/100) = ln(100)
• 23.0, ten times as many classes
• 0.01, one over the number of classes

Show the answer and why

Answer: 4.61, because -ln(1/100) = ln(100)

Cross-entropy charges -log of the probability given to the right class. Spread evenly over C classes that is ln(C): 2.30 for ten, 4.61 for a hundred, and about 10.4 for a 32,000-token vocabulary, which is the number you will see at step 0 of the pretraining lab in Part 12.

Check my answersClear