PyTorch Basics: Tensors, Models and the Training Loop

PyTorch is the most widely used framework for deep learning, especially in research and generative AI. It does the hard parts for you: it calculates gradients automatically, runs on GPUs, and provides ready-made layers, loss functions and optimisers. If you understood the NumPy network in neural networks explained, PyTorch will feel familiar, because it does the same steps with less code.

The building blocks

• Tensors: multi-dimensional arrays, like NumPy arrays, that can live on a GPU and track how they were computed.
• Autograd: PyTorch records the operations on tensors and calculates all gradients with one call, loss.backward(). No hand-written backpropagation.
• nn.Module and nn.Sequential: containers for layers such as nn.Linear (a fully connected layer) and activations such as nn.Tanh or nn.ReLU.
• Loss functions such as nn.BCEWithLogitsLoss (yes/no) or nn.MSELoss (numbers).
• Optimisers such as torch.optim.SGD or torch.optim.Adam, which update the weights.

The training loop

This is the circle network from the first lesson, written in PyTorch. It assumes X and y from that lesson. It was not run in our checking environment, because PyTorch could not be installed there, so no output is shown; run it locally after pip install torch.

import torch
from torch import nn

X_t = torch.tensor(X, dtype=torch.float32)                  # same circle data as before
y_t = torch.tensor(y, dtype=torch.float32).unsqueeze(1)

model = nn.Sequential(nn.Linear(2, 8), nn.Tanh(), nn.Linear(8, 1))   # 2 inputs -> 8 hidden -> 1 output
loss_fn = nn.BCEWithLogitsLoss()                            # cross-entropy for yes/no, applied to raw scores
optimizer = torch.optim.SGD(model.parameters(), lr=1.0)

for epoch in range(2001):
    logits = model(X_t)                                     # forward pass
    loss = loss_fn(logits, y_t)
    optimizer.zero_grad()                                   # clear old gradients
    loss.backward()                                         # backpropagation, done for you
    optimizer.step()                                        # gradient descent step
    if epoch % 500 == 0:
        accuracy = ((logits > 0).float() == y_t).float().mean().item()
        print(f"epoch {epoch:4}: loss {loss.item():.3f}, accuracy {accuracy:.2f}")

Every line maps to a step you have already seen in NumPy:

Step	NumPy version (first lesson)	PyTorch
Forward pass	np.tanh(X @ W1 + b1) and so on	model(X_t)
Loss	cross-entropy formula	loss_fn(logits, y_t)
Backpropagation	hand-written gradient formulas	loss.backward()
Update weights	W1 -= lr * dW1 and so on	optimizer.step()
Reset gradients	not needed	optimizer.zero_grad()

Going further in practice

• Batches: wrap data in a DataLoader to train in mini-batches and shuffle each epoch.
• Train and evaluation modes: call model.train() while training and model.eval() with torch.no_grad() when predicting, so layers like dropout behave correctly and no gradients are stored.
• GPU: move the model and tensors to a device with .to("cuda") when one is available.
• Saving: save the learned weights with torch.save(model.state_dict(), path).
• Pre-trained models: most real projects load a pre-trained model and fine-tune it rather than building from scratch.

Common mistakes

1. Forgetting optimizer.zero_grad(), so gradients pile up across steps.
2. Predicting in training mode, so dropout and batch normalisation behave differently.
3. Mixing devices, with the model on the GPU and data on the CPU.
4. Applying a sigmoid before BCEWithLogitsLoss, which already includes it.

Key takeaway

PyTorch gives you tensors, automatic gradients, layers, losses and optimisers. The training loop is always the same: forward pass, loss, zero_grad, backward, step. Understanding the NumPy version makes each line clear. For a comparison of frameworks, see PyTorch vs TensorFlow; for adapting large models, see fine-tuning and LoRA.