PyTorch 2 Cheatsheet

Tensors

• Install PyTorch — Install the library. Choose your CUDA version at pytorch.org/get-started for GPU builds.
• Import PyTorch — The standard imports for building models.
• Create tensors — Create tensors from data, zeros, ones or random values.
• Check shape, type and device — Inspect a tensor's properties.
• Reshape a tensor — Change the shape. -1 infers one dimension.
• Convert to and from NumPy — Share memory with NumPy arrays on the CPU.
• Multiply matrices — Matrix multiplication with the @ operator.
• Pick the best device — Use a GPU if available, including Apple Silicon.

Autograd and models

• Compute gradients — Track operations and compute derivatives with backpropagation.
• Skip gradient tracking — Save memory and time when predicting.
• Define a model — Subclass nn.Module, create layers in __init__ and define forward.
• Count a model's parameters — Count trainable parameters.
• Load data in batches — Wrap a dataset to get shuffled mini-batches.

Training and saving

• Write a training loop — The standard loop: forward pass, loss, backward pass, update.
• Choose a loss and optimizer — Set up cross-entropy loss and the Adam optimizer.
• Evaluate a model — Switch to evaluation mode and turn off gradients.
• Save a model — Save the weights, which is recommended over pickling the whole model.
• Load a model — Create the model, then load the saved weights.
• Schedule the learning rate — Reduce the learning rate over time. Call scheduler.step() once per epoch.

Performance

• Speed up a model — Compile the model with torch.compile for faster training and inference.
• Use mixed precision — Run operations in lower precision on the GPU to save memory and time.
• Make runs reproducible — Set the random seed for repeatable results.
• Check GPU memory — Print how much GPU memory tensors use.
• Free cached GPU memory — Release unused cached memory back to the GPU.

Related cheatsheets

• Hugging Face cheatsheet — The hf command line, and Transformers and Datasets code for downloading, running, training and sharing models.
• Matplotlib and seaborn 3 cheatsheet — Draw line, scatter, bar and histogram charts with Matplotlib, and statistical charts with seaborn.
• pandas and NumPy 3 cheatsheet — Load, clean, reshape, group and analyze tabular data with pandas, and work with arrays in NumPy.
• Polars 2 cheatsheet — Fast DataFrames in Python: expressions, filtering, grouping, joins, lazy queries and file formats.
• Prisma ORM cheatsheet — Model your database in a schema, run migrations, and query with the type-safe Prisma Client. Examples use the prisma-client-js generator; check the docs for newer generators.
• React 19 cheatsheet — Hooks, components and patterns for React 19 in TypeScript or JavaScript.

Browse all cheatsheets