This code won’t work, as best_model holds a reference to model, which will be updated in each epoch.
You could use copy.deepcopy to apply a deep copy on the parameters or use the save_checkpoint method provided in the ImageNet example 621.
Here is a small example for demonstrating the issue with your code:
model = nn.Linear(10, 2)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

for epoch in range(10):
    optimizer.zero_grad()
    output = model(torch.randn(1, 10))
    loss = criterion(output, torch.randn(1, 2))
    loss.backward()
    optimizer.step()
    
    # Save 2nd epoch
    if epoch == 2:
        best_model = model  # Won't work!
        #best_model = copy.deepcopy(model)  # Will work
        
# Compare models
for param1, param2 in zip(best_model.parameters(), model.parameters()):
    print((param1 == param2).all())