Home › Python Data & ML › PyTorch training CLAUDE.md

PyTorch training CLAUDE.md

CLAUDE.md for PyTorch 2.x model training repos: reproducible configs, data loaders, mixed precision, and checkpoints.

Claude Code (CLAUDE.md)Python Data & ML

How to use

Save this as CLAUDE.md in your repository root, commit it, and your agent will pick it up automatically. Edit the commands to match your project.

# CLAUDE.md: PyTorch training repo

## Commands
- Env: `uv sync` · Tests: `uv run pytest -q` · Lint: `uv run ruff check . && uv run ruff format --check .`
- Smoke run (must work on CPU in < 2 min): `uv run python train.py --config configs/debug.yaml`

## Code rules
- Config lives in YAML/dataclasses (`configs/`); no magic numbers in `train.py`. Every run logs its full resolved config.
- Seed everything (`torch.manual_seed`, `numpy`, `random`) from the config; note that full determinism also needs `torch.use_deterministic_algorithms(True)`.
- Device-agnostic code: pick `device` once (`cuda`, `mps`, `cpu`) and move tensors with `.to(device)`. No hard-coded `.cuda()`.
- Mixed precision via `torch.autocast(device_type, dtype=torch.bfloat16)`; use `GradScaler` only with float16.
- `model.train()` / `model.eval()` and `torch.inference_mode()` for evaluation.
- `DataLoader` with `num_workers`, `pin_memory` on CUDA, and dataset classes that do not load everything into memory.
- `torch.compile` is optional and behind a config flag.
- Checkpoints save model, optimizer, scheduler, step and config with `torch.save`; load with `weights_only=True` when loading only weights.
- Log metrics to the project's tracker (W&B, MLflow or TensorBoard), not print statements.

## Tests
Shape tests for models, a test that a batch overfits in a few steps, and data pipeline tests on a tiny fixture.

## Don't
Don't start long GPU runs, download large datasets, or delete `checkpoints/` without asking.