torch review
torch is the PyPI distribution for PyTorch, a tensor and automatic-differentiation framework used to build, train, and run neural networks on CPUs and several accelerator backends. Its Python API covers tensor operations, nn.Module models, optimizers, data loading, compilation, distributed execution, and serialization. Version 2.13.0 brings FlexAttention to Apple MPS, adds a fused LinearCrossEntropyLoss for large-vocabulary training, introduces the prototype CuTeDSL compiler backend, and expands distributed work. It also removes named tensors and Bazel build files, so this is not a zero-risk minor upgrade.
PyTorch remains the practical default when the surrounding model and tooling ecosystem already assumes it. Do not install 4,560 MB of runtime for classical ML, remote API calls, or a small inference path that can consume an exported model.
We installed it
| Install | ✓ · 31.1s | 29 packages on disk · 4560 MB |
| Import | ✓ | import functorch in 3.04s · compiled extensions · py.typed · requires Python >=3.10 |
| Known vulns | 0 | (pip-audit) |
Answers from our run
Does torch install cleanly?
Yes. In a fresh container with an empty cache, pip install torch finished in 31 seconds, leaving 29 packages and 4560 MB on disk. pip-audit reported no known vulnerabilities.
What does torch need to run?
Python >=3.10, and a platform wheel with compiled extensions. In our run import functorch succeeded in 3.04s, and the package ships py.typed for type checkers.
torch or tensorflow: which should you use?
tensorflow: Use it when an existing production stack is built around Keras, TensorFlow Serving, or TFX. PyTorch remains the practical default when the surrounding model and tooling ecosystem already assumes it.
When should you not use torch?
The problem is tabular regression, trees, clustering, or feature pipelines; scikit-learn is a better fit and avoids a multi-gigabyte install
Use it if
- You train or fine-tune neural networks and the model, checkpoint, or upstream library already targets PyTorch
- GPU tensor operations and automatic differentiation are part of the application rather than an offline preprocessing step
- Your stack uses Hugging Face, torchvision, Lightning, vLLM, or another project centered on torch tensors and modules
- Custom model control, eager debugging, torch.compile, or distributed training justifies a full deep-learning runtime
- The problem is tabular regression, trees, clustering, or feature pipelines; scikit-learn is a better fit and avoids a multi-gigabyte install
- The service only calls a hosted model API; no local tensor framework is needed
- The deployment has a tight image or cold-start budget; our default install occupied 4,560 MB before model weights
- You only serve an exported model and do not need training or autograd; ONNX Runtime can be a narrower inference dependency
- You rely on named tensors or Bazel builds; both are removed in 2.13.0 according to the release notes
Setup reality
Our fresh Python 3.12 install of torch 2.13.0 succeeded in 31.1 seconds. It placed 29 packages on disk using 4,560 MB, and pip-audit found no known vulnerabilities. The distribution declares 17 direct dependencies, requires Python 3.10 or newer, ships compiled .so extensions, and includes py.typed. PyPI does not declare a license value. The measured import probe targeted functorch and completed in 3.04 seconds; it did not time import torch itself.
Wheel choice determines hardware support and much of the footprint. Use the selector on pytorch.org for CPU, CUDA, ROCm, or other supported builds instead of assuming plain pip chose the intended accelerator. A successful import does not prove that the driver, device, and wheel match. Check torch.cuda.is_available(), torch.backends.mps.is_available(), or the relevant backend at application startup, then keep model parameters and input tensors on the same device.
Version 2.13.0 has upgrade traps beyond model code. Named tensor APIs are gone, Bazel build support is removed, and source builds now require newer toolchain pieces including NCCL 2.23 for that configuration. The release notes also track a regression where the ROCm 7.2 wheel can fail on torch.compile CPU work when no GPU is present. Use a standard CPU or CUDA build for GPU-less hosts, or run that ROCm wheel inside the intended ROCm environment.
DataLoader workers are processes. On spawn-based platforms, create them behind an if name == "main" guard and ensure datasets can be serialized. torch.compile pays compilation cost on early calls and can recompile for new shapes or graph breaks. Checkpoints loaded through pickle-capable paths are executable inputs; load only trusted artifacts, prefer state_dict files, use map_location for portability, and review weights_only behavior for the exact saved format.
Patterns
Create a tensor on the selected device create-device-tensor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.tensor([[1.0, 2.0]], device=device)Every tensor participating in an operation must be on a compatible device. Test the backend instead of assuming a CUDA wheel sees a GPU.
Define an nn.Module define-neural-module
class Classifier(torch.nn.Module):
def __init__(self):
super().__init__()
self.layers = torch.nn.Sequential(
torch.nn.Linear(32, 64),
torch.nn.ReLU(),
torch.nn.Linear(64, 4),
)
def forward(self, x):
return self.layers(x)Call model(x), which runs module hooks, instead of invoking forward() directly.
Run a basic training loop train-one-epoch
model.train()
for inputs, targets in loader:
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad(set_to_none=True)
loss = loss_fn(model(inputs), targets)
loss.backward()
optimizer.step()Gradients accumulate by default. Clear them on every ordinary optimization step.
Disable training behavior and gradients run-inference
model.eval()
with torch.inference_mode():
predictions = model(batch.to(device))eval() changes dropout and batch-normalization behavior; inference_mode() disables autograd bookkeeping. Correct inference usually needs both.
Save portable model parameters save-model-state
torch.save({
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": epoch,
}, "checkpoint.pt")Saving state dictionaries is less coupled to Python class locations than pickling the complete model object.
Restore weights on a CPU host load-model-state
checkpoint = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)
model.load_state_dict(checkpoint["model"])
model.eval()Load checkpoints only from trusted sources. Confirm weights_only compatibility with the exact objects saved in your file.
Create a DataLoader batch-data
loader = torch.utils.data.DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=4,
pin_memory=torch.cuda.is_available(),
)Worker processes require a guarded main entrypoint on spawn platforms. More workers can increase memory use and do not always improve throughput.
Compile a model after moving it to its device compile-model
model = model.to(device)
compiled = torch.compile(model)
output = compiled(example.to(device))The first call includes compilation work. Shape changes, unsupported Python, or graph breaks can trigger more compilation or eager fallback.
Use automatic mixed precision on CUDA train-mixed-precision
scaler = torch.amp.GradScaler("cuda")
with torch.amp.autocast("cuda"):
loss = loss_fn(model(inputs), targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()Use torch.amp rather than deprecated torch.cuda.amp spellings. Gradient scaling is used for fp16, not every reduced-precision dtype.
Increase effective batch size by accumulation accumulate-gradients
optimizer.zero_grad(set_to_none=True)
for step, (inputs, targets) in enumerate(loader, 1):
loss = loss_fn(model(inputs), targets) / accumulation_steps
loss.backward()
if step % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)Divide the loss by the number of accumulated steps. Handle the final partial group when the loader length is not divisible.
Train only an added head freeze-backbone
for parameter in model.backbone.parameters():
parameter.requires_grad_(False)
optimizer = torch.optim.AdamW(
(p for p in model.parameters() if p.requires_grad),
lr=1e-4,
)Filtering frozen parameters avoids allocating optimizer state for them. BatchNorm behavior may still need explicit control.
Seed CPU and accelerator generators set-reproducible-seed
torch.manual_seed(20260824)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(20260824)
torch.use_deterministic_algorithms(True)Deterministic algorithms can reduce speed or raise when an operation has no deterministic implementation. DataLoader workers need their own seed plan.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| tensorflow | PyPI | Use it when an existing production stack is built around Keras, TensorFlow Serving, or TFX |
| jax | PyPI | Use it for functional transformations such as jit, grad, and vmap or TPU-focused research |
| mlx | PyPI | Use it for focused array and model work optimized around Apple silicon |
More ai / ml guides
openai · mcp · huggingface-hub · @modelcontextprotocol/sdk · scikit-learn · tiktoken · the whole shelf →
How this guide is made: grounded in the library's documentation, release notes, changelog, and issue history, on a fixed rubric — not a hands-on install of every release. The 50 most-downloaded entries are additionally install-verified in clean containers. Corrections: contact the desk.