How do we modify the early stopping procedure to account for better losses after initial rise in losses?

Question

I have a question regarding the usage of early stopping in the training of my forecasting model. Curious about how the training would go without early stopping, I observed that the test loss seems to converge to a stable value even after an initial uprising (epoch 250 - 400). With early stopping, it is likely that the training would halt within this loss uprising segment, and it will forego long term benefits (such as convergence after epoch 400 for example).

My training block code is as follows:

def training_loop(n_epochs, model, optimizer, loss_fn, train_loader, test_loader, patience=20):
    best_test_loss = float('inf')
    counter = 0  # Counter for early stopping

    for epoch in range(n_epochs):
        
        # Training phase
        model.train()
        train_loss = 0.0
        
        for X_train, y_train in train_loader:
            optimizer.zero_grad()                                                                
            train_preds = model(X_train)  
            # print(train_preds.shape, y_train.shape)
            train_loss_batch = loss_fn(train_preds, y_train)

            train_loss_batch.backward()                                        
            optimizer.step()                                                 
            train_loss += train_loss_batch.item()
        
        training_loss.append(train_loss / len(train_loader))
        
                
        # Evaluation phase
        model.eval()
        test_loss = 0.0
        with torch.no_grad():
            for X_test, y_test in test_loader:
                test_preds = model(X_test)      
                # print(test_preds.shape, y_test.shape)      
                test_loss_batch = loss_fn(test_preds, y_test) 

                test_loss += test_loss_batch.item()

        testing_loss.append(test_loss / len(test_loader))
        
        # Check for early stopping
        if test_loss/len(test_loader) < best_test_loss:
            best_test_loss = test_loss/len(test_loader)
            counter = 0
            print("Found new best test loss, resetting counter")
        else:
            counter += 1
            print("Current best test loss: {:.4f}, counter: {}".format(best_test_loss, counter))
            if counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break