45. Your First Neural Network in PyTorch

45.3 Train Loop

By Ravindra Bagale · Last updated: 24 October 2026

Why. Training is a loop: predict → loss → backward → step.

What. One epoch = one full pass over your training tensors. Beginners can run dozens of epochs on a tiny toy set.

How.

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)
X = torch.tensor([
 [120., 8., 12.],
 [80., 5., 40.],
 [200., 12., 8.],
 [60., 4., 55.],
 [180., 11., 10.],
 [90., 6., 35.],
], dtype=torch.float32)
# fictional reorder labels
y = torch.tensor([[1.], [0.], [1.], [0.], [1.], [0.]])

model = nn.Sequential(nn.Linear(3, 8), nn.ReLU(), nn.Linear(8, 1))
criterion = nn.BCEWithLogitsLoss()
opt = optim.Adam(model.parameters(), lr=0.05)

for epoch in range(40):
 opt.zero_grad()
 logits = model(X)
 loss = criterion(logits, y)
 loss.backward()
 opt.step()
 if epoch % 10 == 0:
 print(epoch, float(loss))

What you should see. Loss trending down (not always monotone every step). This is fictional Mauli-style numbers — not a production trainer.