---

name: pytorch description: >- Assists with building, training, and deploying neural networks using PyTorch. Use when designing architectures for computer vision, NLP, or tabular data, optimizing training with mixed precision and distributed strategies, or exporting models for production inference. Trigger words: pytorch, torch, neural network, deep learning, training loop, cuda. license: Apache-2.0 compatibility: "Requires Python 3.10+ (PyTorch 2.9 and later); a CUDA-capable GPU is recommended" metadata: author: terminal-skills version: "1.2.0" repository: https://github.com/pytorch/pytorch category: data-ai tags: ["pytorch", "deep-learning", "neural-network", "gpu", "machine-learning"]

PyTorch

Overview

PyTorch is a deep learning framework for building and training neural networks with dynamic computation graphs and automatic differentiation. It provides tensor operations with GPU acceleration, nn.Module for defining architectures, DataLoader for efficient data loading, mixed precision training for performance, and export tools (torch.export, ONNX) for production deployment.

Instructions

• Install with pip from the selector at pytorch.org/get-started/locally (it prints the right index URL for your CUDA version or CPU-only). PyTorch 2.9 and later need Python 3.10+.
• Pick the device once: device = "cuda" if torch.cuda.is_available() else "cpu" (Apple Silicon: "mps").
• When defining models, subclass nn.Module with __init__ for layers and forward for computation, using nn.Sequential for simple stacks.
• When training, run the standard loop: optimizer.zero_grad(), forward pass, loss, loss.backward(), clip_grad_norm_, optimizer.step(). Call model.train() before the loop and model.eval() for validation.
• When loading data, subclass Dataset with __len__ and __getitem__, then use DataLoader with num_workers=4, pin_memory=True and persistent_workers=True for GPU training.
• When optimizing performance, use torch.compile(model) (first call is slow while it compiles), mixed precision with torch.autocast(device_type="cuda", dtype=torch.bfloat16), and DistributedDataParallel launched with torchrun for multi-GPU. With float16 also use torch.amp.GradScaler("cuda"); bfloat16 does not need it.
• When doing transfer learning, load pretrained weights with the weights= enum (ResNet50_Weights.DEFAULT; the old pretrained=True was removed) or from Hugging Face, freeze the backbone, and replace the classifier head.
• When saving, store model.state_dict(); load with torch.load(path, map_location=device, weights_only=True). weights_only=True is the default since 2.6 and rejects arbitrary pickled objects.
• When deploying, use torch.export.export(model, example_args) for a graph you can ship, and torch.onnx.export(model, example_args, "model.onnx", dynamo=True) for ONNX (the dynamo exporter is the default since 2.9). TorchScript (torch.jit.trace/script) is deprecated; keep it only for legacy code.
• For INT8 or lower-precision inference, use the torchao package (torchao.quantization.quantize_). torch.ao.quantization is being phased out in favour of it.

Examples

Example 0: Install and verify

pip install torch torchvision
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

Prints the version (for example 2.14.1) and True when a CUDA GPU is usable.

Example 1: Fine-tune a vision model for image classification

User request: "Fine-tune a pretrained ResNet for classifying product images into 12 categories"

import torch, torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights

device = "cuda" if torch.cuda.is_available() else "cpu"
model = resnet50(weights=ResNet50_Weights.DEFAULT)
for p in model.parameters():
    p.requires_grad = False
model.fc = nn.Linear(2048, 12)          # new head is trainable
model.to(device)

opt = torch.optim.AdamW(model.fc.parameters(), lr=1e-3, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
loss_fn = nn.CrossEntropyLoss()

model.train()
for epoch in range(10):
    for images, labels in train_loader:   # DataLoader with RandomCrop/ColorJitter/Normalize
        images, labels = images.to(device), labels.to(device)
        opt.zero_grad()
        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            loss = loss_fn(model(images), labels)
        loss.backward()
        opt.step()
    sched.step()

Result: only the 12-way head is trained, so each epoch is fast; the loss printed per epoch falls steadily and validation accuracy is checked with model.eval() and torch.no_grad().

Example 2: Fine-tune a transformer for sentiment and export it

User request: "Build a sentiment model from bert-base-uncased and export it for serving"

from transformers import AutoTokenizer, AutoModelForSequenceClassification
tok = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
# ... train with AdamW(lr=2e-5), a linear warmup scheduler and clip_grad_norm_(model.parameters(), 1.0)

model.eval()
sample = tok("great battery life", return_tensors="pt")
torch.onnx.export(model, (sample["input_ids"], sample["attention_mask"]),
                  "sentiment.onnx", input_names=["input_ids", "attention_mask"],
                  dynamo=True)

Result: a sentiment.onnx file that ONNX Runtime can load. If export fails on a data-dependent branch, try torch.export.export first to see the exact graph break, then pass dynamic_shapes for variable batch or sequence length.

Guidelines

• Use AdamW over Adam for decoupled weight decay.
• Keep the Dataset on CPU and move batches to the device in the loop.
• Use model.eval() together with torch.no_grad() (or torch.inference_mode()) for inference; forgetting eval() leaves dropout and batch-norm in training mode.
• torch.compile can recompile when input shapes change; mark dynamic dimensions or pad to fixed sizes.
• Never torch.load a checkpoint from an untrusted source with weights_only=False; it can execute arbitrary code.
• Save state_dict, not the whole pickled model, so checkpoints survive code changes.
• Do not mix torch.jit and torch.export in new code; pick torch.export.
• Use PyTorch for custom architectures and research; for plain tabular data, gradient boosting is often faster and just as accurate.