# PyTorch cheat sheet — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-cheat-sheet

> Recall the essential PyTorch APIs and rules on one page, and reproduce a complete training script skeleton from memory.

## Everything on one page

**Tensors**: `shape`, `dtype`, `device`. Models use `float32`; class labels use `int64`. `view` never copies, `reshape` may. `dim=k` in a reduction removes dimension `k`. Broadcasting aligns shapes from the right; watch `(N, 1)` vs `(N,)`.

**Autograd**: `requires_grad=True` records operations; `loss.backward()` fills `.grad` on leaves; gradients **accumulate**. Use `no_grad()` or `inference_mode()` when not training, `.detach()` or `.item()` to cut a value from the graph.

**Models**: subclass `nn.Module`, register layers in `__init__`, compute in `forward`, call `model(x)`. Use `nn.ModuleList`, not a plain list. No softmax before `CrossEntropyLoss`; use `BCEWithLogitsLoss` for binary.

**Loop**: `zero_grad` -> forward -> loss -> `backward` -> `step`. `model.train()` to train, `model.eval()` plus `inference_mode()` to evaluate.

**Data**: `Dataset` (`__len__`, `__getitem__`) -> `DataLoader` (`batch_size`, `shuffle`, `num_workers`). Augment training data only.

**Speed and production**: move model and every batch to `device`; `autocast` for mixed precision; `torch.compile`; save `state_dict`; `torch.load(..., weights_only=True)`.

## The course in five stages

Revise in the same order the course was built: tensors and their shapes, autograd and gradients, models built from nn.Module, the training and evaluation loop, and finally saving, exporting and serving.

![Diagram of 5 connected stages from left to right: tensors, autograd, nn.Module models, training loop, save and deploy.](assets/figures/pytorch/section-10-map.svg) — Figure 10.1: tensors -> autograd -> nn.Module models -> training loop -> save and deploy

## A training script skeleton

If you can write this from memory and explain each line, you understand the core of PyTorch. `train_ds`, `val_ds`, `SmallCNN` and `evaluate` come from earlier topics; swap in your own Dataset and model.

```python
import torch
from torch import nn
from torch.utils.data import DataLoader

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_ds, batch_size=256, shuffle=False)

    model = SmallCNN(num_classes=10).to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    for epoch in range(10):
        model.train()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            loss = loss_fn(model(x), y)
            loss.backward()
            optimizer.step()
        val_loss, val_acc = evaluate(model, val_loader, loss_fn, device)
        print(f"epoch {epoch}: val_loss={val_loss:.3f} val_acc={val_acc:.2%}")

    torch.save(model.state_dict(), "model.pt")

if __name__ == "__main__":
    main()
```

**Quiz:** Which pair is the correct way to evaluate a model without wasting memory and with the right layer behaviour?

- [ ] `model.train()` with `torch.no_grad()`
- [ ] `model.eval()` alone
- [ ] `torch.inference_mode()` alone
- [x] `model.eval()` together with `torch.inference_mode()` (or `no_grad()`)

*Answer:* `model.eval()` together with `torch.inference_mode()` (or `no_grad()`). `eval()` fixes Dropout and BatchNorm behaviour; `inference_mode()` or `no_grad()` stops graph building. You need both.