PyTorch 2 · Training and saving
Write a training loop
The standard loop: forward pass, loss, backward pass, update.
model.train()
for x, y in loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = loss_fn(model(x), y)
loss.backward()
optimizer.step()
Official PyTorch 2 documentation