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.