Since all the real work is done in the train and test loops, the outer “train-test loop” of our example is basically trivial:

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, criterion, optimizer)
    test(test_dataloader, model, criterion)

It’s moderately interesting that there is no off-the-shelf implementation of this; we roll it ourselves in every script.