The model to keep is the one from the epoch with the lowest validation loss, not the last one. Early stopping watches that loss, remembers the best weights, and stops once patience epochs in a row have not beaten it.

Remembering needs a real copy. state_dict() returns the live parameter tensors, which the next optimizer step changes in place:

import copy
import torch
from torch import nn

model = nn.Linear(2, 1)
alias = model.state_dict()
best = copy.deepcopy(model.state_dict())
with torch.no_grad():
    model.weight.fill_(5.0)
print(alias["weight"])  # 5s: it moved
print(best["weight"])  # the old values
model.load_state_dict(best)

To resume training later, save a checkpoint: one dictionary with the model's state, the optimizer's state, which for Adam holds running averages a restart would lose, and the epoch number.

import torch
from torch import nn

model = nn.Linear(2, 1)
opt = torch.optim.Adam(model.parameters())
torch.save({"model": model.state_dict(),
            "optimizer": opt.state_dict(),
            "epoch": 7}, "checkpoint.pt")
saved = torch.load("checkpoint.pt")
opt.load_state_dict(saved["optimizer"])
print(saved["epoch"])

Write early_stopping(model, step, val_loss, patience, max_epochs=100): each epoch call step() then val_loss(), stop after patience epochs in a row without a new lowest loss, load the best epoch's weights back into model, and return that epoch's number, counting from 0. Write save_checkpoint(path, model, optimizer, epoch) and load_checkpoint(path, model, optimizer), which restores both states and returns the epoch.

The explanation opens here once every check passes.