Real datasets do not fit in one step. You train on small batches: each step sees a few dozen examples, and one pass over all of them is an epoch.

import torch
from torch.utils.data import (
    DataLoader, TensorDataset)

x = torch.arange(10.0).unsqueeze(1)
data = TensorDataset(x, x * 2)
loader = DataLoader(
    data, batch_size=4, shuffle=True)
for epoch in range(2):
    for batch_x, batch_y in loader:
        print(epoch, batch_x.flatten())

TensorDataset pairs the rows of x and y; DataLoader cuts them into shuffled batches. On an iPhone or iPad keep the default num_workers=0: an app runs as one process.

Write batches(x, y, size) returning a DataLoader over x and y with that batch size, shuffled. Write train_epochs(model, loader, epochs, lr=0.05) that trains with torch.optim.Adam and nn.MSELoss on every batch of every epoch, and returns the list of average losses per epoch.

The explanation opens here once every check passes.