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

Copy CLAUDE.md

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.

Note: Review and test before use. Provided as is, without warranty. See Terms.

# 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.