A Python Machine Learning Setup You Can Safely Hand to an AI Agent

The hard part of starting a machine learning project in Python is rarely the model. It’s the afternoon you lose to a PyTorch build that can’t see your GPU, a notebook kernel pointed at the wrong interpreter, or a requirements.txt that installs something slightly different on every machine. Add an AI coding assistant to the mix and there’s a new way for it to go wrong, because an agent that hits an ImportError will happily type pip install and move on without telling you which build it grabbed or where it put it.

So this is the setup I’d reach for today: uv to manage Python and the project, PyTorch 2.10 pinned to the exact build you want, and a short instruction file so the agent follows the same rules you do.

Why PyTorch is the awkward dependency

Most Python packages are one wheel per platform on PyPI and that’s the end of it. PyTorch is different, and the uv PyTorch guide explains why pretty clearly. PyTorch publishes a separate build for each accelerator, marks which one you have in the local version (so you’ll see versions like +cpu or +cu128), and hosts many of those builds on its own index at download.pytorch.org instead of PyPI.

For PyTorch 2.10.0, which shipped on January 21, the official install selector offers CUDA 12.6, CUDA 12.8 and CUDA 13.0 builds on Linux, plus ROCm 7.1 for AMD cards. What you get from a plain pip install torch depends on where you run it. On x86 Linux, PyPI serves the CUDA 12.8 build. On Windows, the PyPI package is the CPU build, and on macOS there’s no CUDA build at all.

That Linux default has a cost angle people don’t notice until the bill or the CI timer shows up. The PyTorch 2.10.0 wheel for x86 Linux on PyPI is over 900 MB by itself, and it also declares a long list of NVIDIA runtime packages (cuDNN, cuBLAS, NCCL and friends) as dependencies on that platform. If your test runner has no GPU, all of that is dead weight it downloads on every cold cache.

Start the project with uv

uv (0.10.2 is the current release as I write this) handles the interpreter, the virtual environment and the lockfile in one tool. The standalone installer on macOS and Linux is:

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

Same advice as any script piped into a shell: the uv docs show how to read it first by piping to less instead. Then create a project:

uv init ml-sandbox --python 3.13
cd ml-sandbox

If that Python version isn’t on your machine, uv downloads it for you by default, so there’s no separate pyenv or system Python step. On the version question, PyTorch 2.10 publishes wheels for Python 3.10 through 3.14, and the release notes make Python 3.14 support for torch.compile() a headline feature. The same notes also warn that torch.jit isn’t guaranteed to work on 3.14, so my suggestion is 3.13 unless you specifically want 3.14 and don’t touch TorchScript.

The project ends up with a handful of files that matter. pyproject.toml holds your dependencies, .python-version records the interpreter, and .venv is the environment itself. The important one is uv.lock, which uv creates the first time you run uv run, uv sync or uv lock. It records the exact resolved versions, it’s meant to be committed, and the docs are explicit that it’s managed by uv and shouldn’t be edited by hand.

Pin torch to the build you actually want

Here’s the pattern I like for a project that runs on a laptop and in CI without a GPU, but trains on a CUDA box. It’s adapted from the optional dependencies example in the uv guide, with the version floors bumped to PyTorch 2.10.0 and its matching torchvision 0.25.0:

[project]
name = "ml-sandbox"
version = "0.1.0"
requires-python = ">=3.13"
dependencies = []

[project.optional-dependencies]
cpu = ["torch>=2.10.0", "torchvision>=0.25.0"]
cu128 = ["torch>=2.10.0", "torchvision>=0.25.0"]

[tool.uv]
conflicts = [[{ extra = "cpu" }, { extra = "cu128" }]]

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu128", extra = "cu128" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu128", extra = "cu128" },
]

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

Now uv sync --extra cpu gives you the much smaller CPU build and uv sync --extra cu128 gives you CUDA 12.8. If your driver wants something else, the CUDA 12.6 and 13.0 indexes follow the same URL pattern (/whl/cu126 and /whl/cu130). One catch from the uv docs is that the CUDA extra will fail to install on macOS, since there are no GPU builds for it.

The explicit = true line is doing more than it looks like. It tells uv to use that index only for packages that are pinned to it in tool.uv.sources, so torch and torchvision come from download.pytorch.org while numpy, pillow and everything else still comes from PyPI.

There’s also a gotcha worth knowing about syncing. Extras are never synced unless you ask for them with --extra, and uv sync is exact by default, which means it removes extraneous packages from the environment. Run a bare uv sync after uv sync --extra cpu and torch gets taken back out. uv run does an inexact sync, so it won’t remove anything, but it’s still simplest to always pass the extra you mean.

If you’d rather not deal with any of this for a quick experiment, the uv pip interface has a --torch-backend=auto option that checks your installed CUDA driver and GPUs and picks a matching PyTorch index, falling back to CPU. As of uv 0.10.2 that option only exists in uv pip, not in the project commands above.

Why that index setup is a security control too

This isn’t just tidiness. In December 2022 PyTorch disclosed a compromised nightly dependency: someone uploaded a package called torchtriton to PyPI, the same name as a dependency PyTorch shipped on its own nightly index, and because PyPI took precedence, Linux users who installed nightlies with pip between December 25 and December 30 got the malicious one. It collected system details and files like ~/.ssh/* and sent them out over DNS. Stable builds weren’t affected.

uv’s defaults are built around exactly that failure. Its package index docs describe the default first-index strategy, where once a package is found on one index uv only considers versions from that index, and they cite the torchtriton attack as the reason. The pip-like unsafe-best-match strategy is available, but it has “unsafe” in the name for a reason. This is precisely the kind of setting an agent might flip to make a resolution error go away, which brings us to the next part.

The other one to remember lives in PyTorch itself. Since PyTorch 2.6, torch.load defaults to weights_only=True, which restricts the unpickler to tensors, primitive types, dictionaries and types you explicitly allow. When a checkpoint won’t load under those rules, the error message tells you that setting weights_only=False will likely work but can result in arbitrary code execution. That’s fine for a file you made yourself and a bad idea for something you just downloaded.

Notebooks without the wrong kernel

For Jupyter, the uv docs suggest keeping Jupyter itself out of your project dependencies and adding a kernel instead:

uv add --dev ipykernel
uv run ipython kernel install --user --env VIRTUAL_ENV $(pwd)/.venv --name=ml-sandbox
uv run --with jupyter jupyter lab

Pick the ml-sandbox kernel in the notebook and imports come from the project environment. Inside a notebook, !uv add updates pyproject.toml and the lockfile, while !uv pip install changes the environment without recording anything. That difference matters a lot once an agent is the one typing into cells.

Tell your agent the rules

Coding agents read project instruction files, and that’s the cheapest guardrail you’ll ever add. OpenAI’s Codex reads AGENTS.md files before doing any work, starting at the repository root and walking down to the directory you launched it from, with a combined limit of 32 KiB by default. Claude Code reads a project CLAUDE.md, and its memory docs describe an @path import syntax, so a CLAUDE.md containing just @AGENTS.md pulls in the same file and you only maintain one copy.

What goes in it should be short and specific to the traps above. Something like this:

# AGENTS.md

This project uses uv. Add or remove packages with `uv add` and `uv remove`,
never with pip or `uv pip install`, and never edit uv.lock by hand.

Run code with `uv run --extra cpu ...` on machines without a GPU. Do not run a
bare `uv sync`, because it removes torch.

Do not change `[tool.uv.index]`, `[tool.uv.sources]` or the index strategy.
If a package won't resolve, stop and explain the error instead.

Load checkpoints with `torch.load` defaults. Never pass `weights_only=False`
for a file that did not come from this repository.

Written as plain sentences, it reads just as well to a new teammate as it does to the agent, which is a decent test of whether the rules make sense.

The takeaway

A reproducible ML environment in 2026 comes down to three decisions you make once: which Python, which PyTorch build from which index, and a lockfile everyone installs from. uv lets you write all three into pyproject.toml and uv.lock, and an AGENTS.md gives the AI helping you a real chance of respecting them instead of improvising with pip. Get those in place before the first notebook, and the only thing left to debug is the model.