Guide

PyTorch is the default framework for neural networks in Python. It uses dynamic graphs, a NumPy-like tensor API and an autograd engine that computes the gradients backpropagation needs. JAX and others are good tools, but most tutorials, model releases and interview questions assume PyTorch.

The working vocabulary is small. You need tensors, devices, nn.Module, datasets and the canonical training loop. Once those are familiar, the rest of the API is a lookup.

Here is the core idea in one sentence. A model's weights are tensors, the loss is a single number, and autograd turns that number into a gradient for every weight. The rest of PyTorch is organization around that cycle: where the data lives, how the batches arrive, and how the run is saved.

Concepts in this section

Read them in order. Each one uses the ones before it.

StepConceptWhat it answers
1PyTorch tensorsWhat is a tensor, and why do shapes, devices and views cause bugs?
2PyTorch autogradHow does backward() know the gradients?
3Modules and dataHow do you define a network and feed it batches?
4The training loopIn what order do the steps run, and how do you debug a stalled run?
5Mixed precision and checkpointsHow do you train faster with less memory, and resume a run?

Common mistakes

  • Mixing devices. The model, the batch and the labels must all sit on the same device.
  • Skipping zero_grad or model.eval(). Gradients pile up, and validation runs with dropout on.
  • Trusting shapes you have not printed. Broadcasting can run without an error and still train the wrong function.
  • Saving only the weights. You cannot resume a run or explain it later without the optimizer state and config.
  • Keeping layers in a plain Python list. PyTorch cannot see them, so their weights are never trained or saved. Use nn.ModuleList.

When a run misbehaves, check in this order: the shapes of one batch, the gradients after one backward pass, and then the loss scale. Most silent failures show up in one of those three.

Key takeaways

  • A tensor is an array with a shape, a dtype and a device. Move everything to the same device.
  • Autograd computes gradients for tensors with requires_grad when you call backward().
  • A training loop is zero_grad, forward, loss, backward and step. Validate in eval mode with inference_mode().
  • Overfit one batch before a long run. Debug shapes, then gradient flow, then the loss scale.
  • Use bfloat16 autocast when the hardware supports it. Save the config with every checkpoint.

Further learning

The official walkthrough of tensors, modules and training is a good companion:

PyTorch: Learn the BasicsLearn the Basics — PyTorch Tutorials 2.14.0+cu130 documentationpytorch.org

Karpathy's Zero to Hero course builds intuition first, then the PyTorch habits:

Neural Networks: Zero to HeroNeural Networks: Zero To Herokarpathy.ai