PyTorch Basics

10 examples to get you started with Deep Learning - 7 basic and 3 intermediate.

Prerequisites

uv venv && source .venv/bin/activate
uv pip install "torch>=2.6" "torchvision>=0.21"
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

• Python 3.14.0 with PyTorch 2.6+.
• CUDA GPU optional but recommended for training (CPU works for small examples).

Basic Examples

1. Create a Tensor

Tensors are PyTorch's core data structure - like NumPy arrays with GPU support.

import torch
 
x = torch.tensor([1.0, 2.0, 3.0])
matrix = torch.zeros(3, 4)
rand = torch.randn(2, 3, generator=torch.Generator().manual_seed(42))
print(x.shape, x.dtype, x.device)

• torch.tensor creates from Python lists; torch.zeros/randn allocate memory.
• .shape, .dtype, and .device describe tensor metadata.
• Set a generator seed for reproducible random tensors.

Related: Autograd - tensors track gradients

2. Tensor Operations

Element-wise and matrix operations work like NumPy.

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])
print(a + b)
print(a @ b.T)  # matrix multiply
print(a.mean(dim=0))

• @ is matrix multiplication; * is element-wise.
• dim=0 reduces along rows (column means).
• Broadcasting applies smaller tensors across larger ones automatically.

3. Move Tensors to GPU

Accelerate computation by placing tensors on CUDA devices.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.randn(1000, 1000, device=device)
y = x @ x.T
print(y.device)

• Always create or move tensors to the target device before operations.
• Mixing CPU and GPU tensors raises errors - keep everything on one device.
• torch.cuda.is_available() returns False on CPU-only machines.

Related: GPUs & Mixed Precision - CUDA and autocast

4. Convert NumPy and PyTorch

Bridge between ecosystems with zero-copy when possible.

import numpy as np
 
np_arr = np.array([1.0, 2.0, 3.0])
tensor = torch.from_numpy(np_arr)
back = tensor.numpy()

• from_numpy shares memory with the NumPy array (in-place changes affect both).
• .numpy() requires CPU tensor and may copy if gradients are attached.
• Use .detach().cpu().numpy() for tensors with requires_grad=True.

5. Indexing and Reshaping

Select, slice, and reshape tensor dimensions.

x = torch.arange(12).reshape(3, 4)
print(x[1, 2])       # single element
print(x[:, 1])       # column
print(x.view(2, 6))  # reshape (contiguous)
print(x.unsqueeze(0).shape)  # add batch dim: (1, 3, 4)

• reshape and view change shape; view requires contiguous memory.
• unsqueeze/squeeze add/remove size-1 dimensions.
• Batch dimension is typically first: (batch, channels, height, width).

6. Enable Gradient Tracking

Tensors that need gradients set requires_grad=True.

w = torch.tensor([2.0, 3.0], requires_grad=True)
loss = (w ** 2).sum()
loss.backward()
print(w.grad)  # tensor([4., 6.])

• .backward() computes gradients via autograd.
• .grad holds the gradient after backward.
• Wrap parameters in nn.Parameter for model training.

Related: Autograd - computation graphs

7. Build a Layer with nn.Module

torch.nn provides differentiable building blocks.

import torch.nn as nn
 
layer = nn.Linear(10, 5)
x = torch.randn(32, 10)  # batch of 32
out = layer(x)
print(out.shape)  # (32, 5)
print(layer.weight.shape)  # (5, 10)

• nn.Linear applies y = xW^T + b.
• First dimension is always batch size.
• Layers are callable: layer(x) runs the forward computation.

Related: Building Models with nn.Module - full models

Intermediate Examples

8. Simple Training Step

Manual forward, loss, backward, and optimizer step.

import torch.nn as nn
import torch.optim as optim
 
model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()
 
x = torch.randn(32, 10)
y = torch.randn(32, 1)
 
optimizer.zero_grad()
pred = model(x)
loss = loss_fn(pred, y)
loss.backward()
optimizer.step()
print(f"loss: {loss.item():.4f}")

• zero_grad() clears previous gradients before each step.
• loss.item() extracts a Python float from a scalar tensor.
• This is the inner loop of every training run.

Related: Training Loops - full epoch loops

9. Save and Load Tensor Checkpoints

Persist model state for later use.

import torch.nn as nn
 
model = nn.Linear(10, 1)
torch.save(model.state_dict(), "linear.pt")
loaded = nn.Linear(10, 1)
loaded.load_state_dict(torch.load("linear.pt", weights_only=True))
loaded.eval()

• state_dict() saves learnable parameters only.
• weights_only=True (PyTorch 2.6+) prevents unsafe pickle loading.
• Call .eval() before inference to disable dropout behavior.

Related: Saving, Loading & Export - full checkpoint patterns

10. Compile a Model with torch.compile

PyTorch 2.x can JIT-optimize models for faster execution.

import torch
import torch.nn as nn
 
model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
model = torch.compile(model)
x = torch.randn(64, 784)
out = model(x)  # first call compiles; subsequent calls are faster
print(out.shape)

• torch.compile fuses operations via TorchInductor.
• First forward pass has compilation overhead.
• Best gains on GPU with repeated inference or training loops.

Stack versions: This page was written for Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, and uv 0.6+. 636f64656775696465732e696f7c6367696f3735347c323032363037