PyTorch

Overview

PyTorch is an open-source framework for building and training neural networks in Python. It gives you tensors, which are multi-dimensional arrays that can live on a CPU or a GPU. It records the operations you perform on them so it can compute gradients automatically. On top of that it supplies layers, loss functions, optimizers, data loading utilities, a compiler and tools for training across many GPUs.

The defining trait of PyTorch is that it runs your code as you write it. Each line executes straight away, and you can print a tensor, step through a model in a debugger or use an ordinary Python if statement inside a network. This style is called eager execution. It made PyTorch popular with researchers first, and industry followed once the same code could be compiled and deployed.

A helpful picture is cooking with a recipe you can taste as you go. Some kitchens make you write the whole recipe, seal it, and only then see what comes out of the oven. PyTorch lets you add an ingredient, taste, adjust and carry on. When the dish is right, you can still hand the recipe to a faster professional kitchen, which is what torch.compile does.

What it is

PyTorch is a library with several layers, from raw arrays up to distributed training.

The pieces fit together in a standard pattern. A model is a subclass of nn.Module that holds parameters. Data arrives through a DataLoader. A loss compares predictions with targets. Autograd turns that loss into gradients, and an optimizer uses the gradients to update the parameters.

PyTorch also has a wide ecosystem around it: domain libraries such as TorchVision and TorchAudio, higher-level training libraries such as PyTorch Lightning, and model hubs such as Hugging Face, whose Transformers library now uses PyTorch as its only backend.

Where it came from

PyTorch grew out of Torch, an older machine learning library with a Lua front end. A team at Facebook AI Research, with Adam Paszke, Sam Gross, Soumith Chintala and Gregory Chanan as original authors, built a Python-first successor and released it in September 2016. It was released under a modified BSD licence.

Several milestones shaped it:

• 2018: Caffe2, Facebook's production-oriented framework, was merged into PyTorch, bringing research and deployment code under one roof.
• September 2022: Meta announced that PyTorch would be governed by the new PyTorch Foundation, part of the Linux Foundation. This moved control from a single company to a neutral body with several member organisations.
• March 2023: PyTorch 2.0 was released. Its headline feature was torch.compile, built on TorchDynamo, a component that captures Python bytecode into graphs without forcing users to rewrite their models.
• 2025: The PyTorch Foundation became an umbrella for related projects, welcoming vLLM and DeepSpeed as hosted projects and later Ray.

Releases continue on a regular cadence within the 2.x line. Version 2.15 was announced at the end of October 2026.

How it works

Tensors and devices

Every computation starts with tensors. A tensor has a shape, a data type such as float32 or bfloat16, and a device. Moving a model or tensor to a GPU is a call like .to("cuda"). Operations on GPU tensors launch GPU kernels, usually through CUDA libraries.

Building the graph as you go

When a tensor has requires_grad=True, every operation on it is recorded. PyTorch builds a graph of these operations dynamically during the forward pass. Because the graph is rebuilt on every run, it can change from one batch to the next. A loop that runs a different number of times per input is no problem.

Backward pass

Calling loss.backward() walks that recorded graph in reverse. It applies the chain rule at each step and adds the result to each parameter's .grad field. This is backpropagation handled for you. Gradients accumulate by default, which is why training loops call optimizer.zero_grad() before each step.

The training loop

A typical step is short:

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

Compilation

torch.compile(model) returns a compiled version of the model. TorchDynamo traces the Python code into graphs, and a back end (by default TorchInductor) fuses operations and generates efficient kernels, often using Triton on GPUs. When it meets Python it cannot capture, it falls back to eager execution for that part, which is called a graph break.

Scaling out

For more than one GPU, DistributedDataParallel copies the model to each device and averages gradients. FSDP shards parameters, gradients and optimizer state across devices so models larger than one GPU's memory can train. Communication between GPUs typically runs through NCCL.

Key concepts and terms

A worked example

This is an illustrative scenario. A small team wants to classify photos of manufactured parts as "good" or "defective". They have 20,000 labelled images at 224 by 224 pixels.

1. They write a Dataset that reads each image and its label, and wrap it in a DataLoader with a batch size of 64 and four worker processes.
2. They load a pretrained ResNet from TorchVision and replace its final layer with a two-output Linear layer. This is transfer learning.
3. They choose cross-entropy loss and AdamW with a learning rate of 0.0003.
4. They train for ten epochs. With 20,000 images and a batch of 64, each epoch is about 313 steps.
5. They wrap the model with torch.compile and turn on mixed precision with torch.autocast. Each epoch gets faster, and memory use drops.
6. After each epoch they evaluate on a held-out set of 2,000 images and save the state_dict when validation accuracy improves.
7. For serving they export the model with torch.export or to ONNX and run it with ONNX Runtime.

The important lesson is that every step is plain Python they can inspect. When accuracy stalls, they can print predictions for a single batch and look.

Strengths and limitations

Strengths:

• Code reads like ordinary Python and is easy to debug.
• Dynamic graphs handle variable-length inputs and control flow naturally.
• A very large ecosystem of models, libraries and tutorials.
• Strong GPU support, plus support for other accelerators.
• torch.compile gives speed-ups without rewriting models.
• Neutral governance through the PyTorch Foundation.

Limitations:

• Deployment outside Python needs an extra export step and separate runtimes.
• Compilation can hit graph breaks that reduce the benefit.
• Large-scale distributed training still needs careful configuration.
• Freedom to write any loop means teams must impose their own structure.

How it compares with alternatives

Practical guidance and best practices

1. Start in eager mode and get the model correct before optimising anything.
2. Add torch.compile once the model is stable, and check for graph breaks if speed-ups are smaller than expected.
3. Use mixed precision, usually bfloat16 on recent GPUs, to save memory and time.
4. Keep data loading off the critical path with several workers and pinned memory.
5. Save checkpoints as state_dict files, not whole pickled models, so code changes do not break loading.
6. Set random seeds and record PyTorch, CUDA and library versions with each run.
7. Use DDP for models that fit on one GPU and FSDP when they do not.
8. Use model.eval() and torch.no_grad() or torch.inference_mode() during evaluation.

Common mistakes

• Forgetting optimizer.zero_grad(), so gradients from previous steps keep adding up.
• Leaving the model in training mode during evaluation, which keeps dropout active and changes normalisation behaviour.
• Mixing tensors on different devices and getting device mismatch errors.
• Calling .item() or printing GPU tensors inside the hot loop, which forces slow synchronisation.
• Loading untrusted pickled checkpoints, which can execute arbitrary code.
• Measuring speed without warming up first, so compilation time distorts the numbers.

Frequently asked questions

Is PyTorch better than TensorFlow?

Neither is better in every situation. PyTorch dominates research and most new model releases, and many production systems use it. TensorFlow keeps a strong position in established production pipelines and on-device deployment. Pick the one that matches your team's skills, your models and your deployment target.

What does torch.compile actually do?

It captures your model's Python code into graphs, optimises them by fusing operations and generates faster kernels for your hardware. You call it once on a model or function, and the first run is slower because of compilation. Later runs are usually faster.

Do I need a GPU to use PyTorch?

No. PyTorch runs on CPUs, and small models train fine there. A GPU or other accelerator becomes important for large models, images, audio and language tasks, where training on a CPU would take far too long.

What is the difference between DDP and FSDP?

DDP gives every GPU a full copy of the model and averages gradients between them. FSDP splits the model's parameters, gradients and optimizer state across GPUs and gathers pieces only when needed. FSDP uses less memory per GPU, so it can train much larger models.

Who controls PyTorch?

Since 2022 PyTorch has been governed by the PyTorch Foundation, part of the Linux Foundation. Meta remains a major contributor, alongside many other companies and individuals.

How do I deploy a PyTorch model?

Common routes are torch.export for an ahead-of-time graph, ONNX export for runtimes such as ONNX Runtime or TensorRT, and serving frameworks that load PyTorch models directly. For large language models, dedicated engines such as vLLM are common.

Glossary

• Autograd: the engine that records tensor operations and computes gradients.
• Backward pass: the reverse walk through the graph that produces gradients.
• Checkpoint: a saved copy of model and optimizer state.
• CUDA: NVIDIA's GPU programming platform that PyTorch uses for GPU work.
• Eager mode: running each operation immediately rather than building a graph first.
• Kernel: a small program that runs one operation on an accelerator.
• Module: a reusable building block that holds parameters and defines a forward pass.
• Optimizer: the algorithm that updates parameters from gradients.
• Sharding: splitting tensors across devices so each holds only a part.
• TorchDynamo: the component of torch.compile that captures Python code into graphs.

Summary

PyTorch is an open-source deep learning framework built around tensors, automatic differentiation and eager, Pythonic execution. It began at Facebook AI Research in 2016, moved to neutral governance under the PyTorch Foundation in 2022, and added torch.compile with version 2.0 in 2023. It suits everything from a first experiment on a laptop to sharded training across many GPUs, provided teams add structure, compile thoughtfully and plan their export and serving path early.

Related

It belongs to Frameworks in my stack, under Deep learning. Others in the same group: