Remember spelling out the chain rule by hand in the NumPy posts just to train one tiny network? This PyTorch tutorial hands you the calculator instead: tensors that run on a GPU, autograd that computes gradients for you, and the one training loop every project reuses, from a 50-line classifier to a giant language model.

“PyTorch is Python first. It feels like NumPy with GPU support and automatic differentiation built in. That is why researchers love it.”

Soumith Chintala, PyTorch creator

Last Updated: July 2026 | Tested on: Python 3.14.6, PyTorch 2.12.1 | Difficulty: Intermediate | Reading Time: 13 minutes

In the earlier posts we built neural networks by hand with NumPy. That was great for understanding the math, but nobody ships production networks that way. Writing the chain rule out by hand is like doing long division when you own a calculator. It works, but why would you? PyTorch is the calculator. It gives you tensors that behave like NumPy arrays but can run on a GPU (Graphics Processing Unit), automatic differentiation that computes gradients for you, and ready-made building blocks (nn.Module, DataLoader, optimizers) so you stop writing the same plumbing over and over. PyTorch is the most common framework in deep learning research today, and tools like torch.compile have made it a solid choice for production too.

The mental model is short enough to memorize. Define your network as a class, pour your data into a DataLoader, write a training loop that does forward, backward, update, then check accuracy on a test set. Every PyTorch project follows this same shape, whether it is a 50-line classifier on your laptop or a billion-parameter language model on a cluster. Learn it once here and you will recognize it everywhere.

Here is what this PyTorch tutorial covers:

• PyTorch tensors: creation, operations, GPU acceleration
• Autograd: automatic gradient computation
• nn.Module: defining network architectures
• DataLoader: batching and shuffling data
• The complete train-evaluate loop

Prerequisites

The diagram shows the PyTorch training cycle. Data goes in as tensors, the model produces a prediction, the loss function measures how wrong it was, .backward() works out which way to nudge every weight, the optimizer takes the step, and you zero the gradients before the next batch. PyTorch records the operations as you run them, so the graph is rebuilt on every forward pass. That “define-by-run” style is why you can use plain Python if statements and loops inside a model and debug it with a normal breakpoint, exactly the reason researchers reach for PyTorch first.

Tensors: NumPy on Steroids

Every PyTorch tutorial starts with tensors, and for good reason. A tensor is a multi-dimensional array, the same idea as a NumPy ndarray, with two extra powers: it can live on a GPU so thousands of numbers get crunched in parallel, and it can remember the operations done to it so gradients fall out automatically. Think of a tensor as a spreadsheet that also keeps a history of every formula you typed, so it can later tell you how each cell affected the final total. In deep learning everything is a tensor: your input data, every weight, every gradient.

📄 pytorch_tensors.py: tensor basics and autograd

import torch
import numpy as np

# Creating tensors
x = torch.tensor([1.0, 2.0, 3.0])
zeros = torch.zeros(3, 4)
randn = torch.randn(2, 3)  # Normal distribution

# From NumPy (shared memory!)
np_array = np.array([1.0, 2.0, 3.0])
tensor_from_np = torch.from_numpy(np_array)

# Aditi explores tensor operations
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

print(f"Shape: {a.shape}")
print(f"Element-wise multiply:\n{a * b}")
print(f"Matrix multiply:\n{a @ b}")
print(f"Sum: {a.sum()}, Mean: {a.mean()}")

# GPU check
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"\nUsing device: {device}")

# Autograd: automatic differentiation
w = torch.tensor([2.0], requires_grad=True)
x = torch.tensor([3.0])
y = w * x     # y = 2 * 3 = 6
loss = (y - 5) ** 2  # (6 - 5)^2 = 1

loss.backward()  # Compute gradient
print(f"\nAutograd demo:")
print(f"  w={w.item()}, x={x.item()}, y={y.item()}, loss={loss.item()}")
print(f"  dLoss/dw = {w.grad.item()}")  # d/dw (wx-5)^2 = 2(wx-5)*x = 2*1*3 = 6

▶ Output

Shape: torch.Size([2, 2])
Element-wise multiply:
tensor([[ 5., 12.],
        [21., 32.]])
Matrix multiply:
tensor([[19., 22.],
        [43., 50.]])
Sum: 10.0, Mean: 2.5

Using device: cpu

Autograd demo:
  w=2.0, x=3.0, y=6.0, loss=1.0
  dLoss/dw = 6.0

What happened here: For everyday math, tensors behave just like NumPy arrays. The new trick is autograd. The moment we set requires_grad=True on w, PyTorch starts writing down every operation that touches it. When we call loss.backward(), it replays that history in reverse and computes dLoss/dw = 2*(wx-5)*x = 2*1*3 = 6. That is the exact same chain-rule math we coded by hand in the backpropagation tutorial, except here it took one line instead of a page. One thing worth knowing: torch.from_numpy shares memory with the original array, so changing the NumPy array also changes the tensor, and the other way around.

Building a Network with nn.Module

📄 pytorch_network.py: defining a neural network class

import torch
import torch.nn as nn

torch.manual_seed(42)  # so your numbers match the output below

class IrisClassifier(nn.Module):
    def __init__(self, input_size=4, hidden_size=32, num_classes=3):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 16),
            nn.ReLU(),
            nn.Linear(16, num_classes),
        )

    def forward(self, x):
        return self.network(x)

# Anvay creates and inspects the model
model = IrisClassifier()
print(model)

total_params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"\nTotal parameters: {total_params:,}")
print(f"Trainable: {trainable:,}")

# Test forward pass with random data
x_test = torch.randn(5, 4)
output = model(x_test)
print(f"\nInput shape: {x_test.shape}")
print(f"Output shape: {output.shape}")
print(f"Output (logits):\n{output}")

▶ Output

IrisClassifier(
  (network): Sequential(
    (0): Linear(in_features=4, out_features=32, bias=True)
    (1): ReLU()
    (2): Linear(in_features=32, out_features=16, bias=True)
    (3): ReLU()
    (4): Linear(in_features=16, out_features=3, bias=True)
  )
)

Total parameters: 739
Trainable: 739

Input shape: torch.Size([5, 4])
Output shape: torch.Size([5, 3])
Output (logits):
tensor([[-0.1325, -0.1591,  0.1821],
        [-0.1557, -0.1128,  0.3271],
        [-0.2509, -0.1991,  0.1356],
        [-0.2263, -0.2010,  0.1915],
        [-0.0727, -0.1913,  0.1214]], grad_fn=<AddmmBackward0>)

What happened here: nn.Sequential stacks layers in order, like an assembly line where each station hands its result to the next. Every nn.Linear is a fully connected layer, just a weight matrix plus a bias. Count them up and the model has 739 trainable numbers: (4*32+32) + (32*16+16) + (16*3+3) = 160 + 528 + 51 = 739. The output is raw scores called logits, not probabilities, because PyTorch’s CrossEntropyLoss applies softmax for you later.

The grad_fn tag on the output is PyTorch telling you it is still tracking operations, ready to compute gradients when you call backward. Because we set torch.manual_seed(42) at the top, your logits will match these exactly; drop the seed and you will get different random weights every run.

The Complete Training Loop

Think of training like teaching a friend to cook a new dish by taste. They make a small batch, you tell them how far off the flavour is, they adjust the spices, and they try again. Do that enough times and the taste locks in. The training loop is exactly this in code: the model makes a prediction, the loss says how wrong it was, the gradients say which way to adjust, and the optimizer makes the adjustment, over and over.

📄 pytorch_train.py: full training pipeline

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

torch.manual_seed(42)  # reproducible weights and shuffling

# Load and prepare data
iris = load_iris()
X = iris.data.astype("float32")
y = iris.target.astype("int64")

scaler = StandardScaler()
X = scaler.fit_transform(X)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

train_ds = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
test_ds = TensorDataset(torch.tensor(X_test), torch.tensor(y_test))
train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=16)

# Viraj trains the classifier
model = IrisClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(51):
    model.train()
    total_loss = 0
    for X_batch, y_batch in train_loader:
        optimizer.zero_grad()
        output = model(X_batch)
        loss = criterion(output, y_batch)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    if epoch % 10 == 0:
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for X_batch, y_batch in test_loader:
                preds = model(X_batch).argmax(dim=1)
                correct += (preds == y_batch).sum().item()
                total += len(y_batch)
        acc = correct / total
        print(f"Epoch {epoch:>3d} | Loss: {total_loss/len(train_loader):.4f} | Test Acc: {acc:.1%}")

▶ Output

Epoch   0 | Loss: 0.8863 | Test Acc: 80.0%
Epoch  10 | Loss: 0.0651 | Test Acc: 100.0%
Epoch  20 | Loss: 0.0496 | Test Acc: 100.0%
Epoch  30 | Loss: 0.0593 | Test Acc: 100.0%
Epoch  40 | Loss: 0.0447 | Test Acc: 100.0%
Epoch  50 | Loss: 0.0609 | Test Acc: 100.0%

What happened here: This is the training loop you will reuse in every project long after this PyTorch tutorial, so it is worth burning into memory. Five steps repeat for each batch: zero_grad() wipes the previous gradients, the forward pass produces an output, criterion measures the loss, backward() computes the gradients, and step() nudges the weights. Iris is a tiny, easy dataset, so the model starts at 80% accuracy after a single pass and locks in at 100% by epoch 10; after that the loss just wobbles around a low value as the weights settle.

The model.eval() and torch.no_grad() calls switch the model into evaluation mode and skip gradient tracking, which saves memory and keeps the test pass honest. Your exact loss numbers will track these closely because of the seed, but a different PyTorch build can shift the last decimal or two.

Saving and Loading Models

📄 save_load.py: persisting trained models

import torch

# Save model weights (recommended approach)
torch.save(model.state_dict(), "iris_model.pth")
print("Model saved to iris_model.pth")

# Load model
loaded_model = IrisClassifier()
loaded_model.load_state_dict(torch.load("iris_model.pth", weights_only=True))
loaded_model.eval()

# Verify loaded model works
sample = torch.tensor(X_test[:3], dtype=torch.float32)
with torch.no_grad():
    preds = loaded_model(sample).argmax(dim=1)
print(f"Predictions: {preds.tolist()}")
print(f"Actual:      {y_test[:3].tolist()}")

▶ Output

Model saved to iris_model.pth
Predictions: [1, 0, 2]
Actual:      [1, 0, 2]

What happened here: state_dict() saves only the learned weights, not the Python class that defines the architecture. Think of it like saving just the filled-in answers to a form rather than the whole form: smaller, more portable, and the approach the PyTorch team recommends. The catch is that you must rebuild the same model class first, then pour the weights back into it, which is why we create a fresh IrisClassifier() before loading. The weights_only=True flag tells PyTorch to load plain tensors only and refuse to run any code hidden in the file, which blocks a known pickle-based attack. In PyTorch 2.12.1 that is already the default, but writing it out makes your intent obvious to the next reader.

Common Mistakes

Interview Corner

Q: What is the difference between torch.save(model) and torch.save(model.state_dict())?

torch.save(model) pickles the whole object, including a reference to the class. That feels convenient until you rename or move the class, and the file refuses to load. torch.save(model.state_dict()) saves only the weight tensors as a dictionary, so it is smaller, survives refactors, and is the approach the PyTorch docs recommend. Save the weights, rebuild the class, load the weights back in.

Practice Exercises

1. Modify the IrisClassifier to use 3 hidden layers with dropout. Compare accuracy with the original.
2. Replace Adam with SGD (Stochastic Gradient Descent) optimizer and compare training curves.
3. Load MNIST from torchvision and train a classifier to 98%+ accuracy.
4. Move the model and data to GPU (if available) and benchmark training speed difference.

More in this series:

• DL: Activation Functions: sigmoid, tanh, ReLU, softmax
• DL: PyTorch GPU Training, Debugging, and Speedups
• Python: PyTorch vs TensorFlow vs JAX, Deep Learning Frameworks Compared

Frequently Asked Questions

Interview Questions on PyTorch

How interviewers actually probe this topic: real scenarios, with answers you can say out loud.

Q: Why do you call optimizer.zero_grad() before backward() on every step?

PyTorch accumulates gradients into the .grad attribute by default instead of overwriting them. If you skip zero_grad(), each batch adds its gradients on top of the last batch’s, so your updates use stale, piled-up values and training becomes erratic. Zeroing first gives every batch a clean slate. The accumulation behaviour is actually useful when you deliberately want it, for example simulating a larger batch size by summing gradients over several small batches before stepping.

Q: What does requires_grad=True do, and which tensors need it?

Setting requires_grad=True tells autograd to track every operation on that tensor so it can compute gradients during backward(). You need it on the parameters you want to learn, which is why nn.Module layers turn it on for their weights automatically. You do not want it on raw input data or on anything inside a torch.no_grad() block, since tracking there just wastes memory building a graph you never use.

Q: Why does nn.CrossEntropyLoss expect raw logits instead of softmax probabilities?

CrossEntropyLoss combines log_softmax and negative log likelihood in one step, which is more numerically stable than doing softmax yourself and then taking a log. If you apply softmax first and pass those probabilities in, you double-apply the squashing, the gradients get muddy, and the model trains poorly. So the last layer should output plain scores, and the loss function handles the rest.

Q: Scenario: your model hits 100% training accuracy, but at inference on the very same data the predictions swing around and look worse. What do you check first?

Check that you called model.eval() before inference. Layers like dropout and batch normalization behave differently in training and evaluation mode: dropout randomly zeros activations during training, and batchnorm uses running statistics only in eval mode. If you forget to switch, dropout is still firing and predictions become noisy and inconsistent. Wrap the inference in torch.no_grad() too, then switch back to model.train() before the next training pass.

Q: Scenario: training runs fine, but GPU memory keeps climbing across epochs until you hit an out-of-memory crash. What is the likely cause?

The classic culprit is holding onto tensors that still carry their computation graph. If you accumulate the loss with total_loss += loss instead of total_loss += loss.item(), every batch’s graph stays alive and memory grows without bound. The same happens if you run evaluation without torch.no_grad(), so the graph is built and retained even though you never call backward. Use .item() or .detach() for values you only want to log, and wrap all inference in no_grad().

Q: How is a PyTorch tensor different from a NumPy array?

For plain math they behave almost identically, and you can convert between them cheaply with torch.from_numpy and .numpy(). The two extra powers a tensor has are that it can live on a GPU for massively parallel computation, and it can track operations through autograd so gradients are computed automatically. One catch worth mentioning: torch.from_numpy shares the underlying memory, so editing the NumPy array also changes the tensor.

What’s Next?

In this PyTorch tutorial you built and trained a real neural network, and you now know the whole loop end to end: tensors and autograd, defining a model with nn.Module, batching data with DataLoader, the train-evaluate cycle, and saving and loading weights. That same shape scales from this tiny Iris classifier all the way up to giant language models, so you already have the pattern you will reuse everywhere. Next, in the TensorFlow and Keras tutorial, we rebuild this same network with Keras so you can feel the difference between the two styles, and in the PyTorch vs TensorFlow comparison we put them head to head to help you pick one for your own projects.

Want the full roadmap from Python basics through deep learning? Head back to the Python + AI/ML tutorial series home to see where this post fits and what to read next.

Go deeper: when you outgrow this post, PyTorch documentation is the next stop.

Related Posts

Previous: Backpropagation Explained: Gradient Descent in Python Made Simple

Next: PyTorch Dataset and DataLoader: The Real Training Loop

Series Home: Python + AI/ML Tutorial Series