def train_secondary_learner(
    secondary_learner, train_dataset, max_epochs, batch_size, eval_freq=50, igf_model_path="secondary_learner.pt"
):
    """
    Train the secondary learner (igf_model)

    Args:
        secondary_learner: secondary learner
        train_dataset: data to train secondary learner
        max_epochs: number of epochs to train secondary learner
        batch_size: batch size of training data of secondary learner
        eval_freq: secondary model evaluation can be triggered at eval_freq
        igf_model_path: path to store trained secondary learner

    Returns:
        Trained secondary learner

    """
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # We will use the first 512 pairs from our dataset as a test set for
    # our secondary learner and the rest to train
    test_dataset = train_dataset[:512]
    train_dataset = train_dataset[512:]
    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
    test_dataloader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)

    # secondary learner model set up
    loss = nn.MSELoss()
    test_loss = nn.MSELoss(reduction="sum")
    secondary_learner.to(device)
    q_optimizer = torch.optim.Adam(secondary_learner.parameters(), lr=0.00001)
    secondary_learner.train()

    # TODO in original code this is written as number of actual batches seen
    # not number of items seen but other places it is number of items instead.
    # improve consistency! changed this to epochs for clarity
    best_test_loss = float("inf")
    # Iterate through batches until we've used max_steps batches
    for epoch in range(int(max_epochs)):
        tr_q_loss = 0.0
        secondary_learner.train()
        for step, batch in enumerate(train_dataloader):
            context = batch[0].to(device)
            real_q = batch[1].to(device)
            predicted_q = secondary_learner(context)
            q_optimizer.zero_grad()
            q_loss = loss(predicted_q, real_q.float())
            q_loss.backward()
            q_optimizer.step()
            tr_q_loss += q_loss.item()

            # model trains fairly quickly so we won't wait for a full epoch
            # eval is triggered at eval_freq and end of epochs
            if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(train_dataloader)):
                tr_loss = tr_q_loss / (step + 1)

                secondary_learner.eval()
                q_loss2 = 0.0
                sum_q2 = 0.0
                predicted = []
                actual = []
                # Compute performance of the secondary learner after this batch
                for step2, batch2 in enumerate(test_dataloader):
                    features2 = batch2[0].to(device)
                    real_q2 = batch2[1].to(device)
                    predicted_q2 = secondary_learner(features2)
                    q_loss2 += test_loss(predicted_q2, real_q2).item()
                    sum_q2 += torch.sum(predicted_q2).item()
                    for ei, i in enumerate(predicted_q2.cpu().detach().numpy()):
                        predicted.append(i.item())
                    for ei, i in enumerate(real_q2.cpu().detach().numpy()):
                        actual.append(i.item())

                q_loss2 /= len(test_dataset)
                print(
                    "Epoch: ",
                    epoch,
                    "step: ",
                    step,
                    "Avg. q:",
                    sum_q2 / len(test_dataset),
                    "Train Loss: ",
                    tr_loss,
                    "Test Loss: ",
                    q_loss2,
                )
                if q_loss2 < best_test_loss:
                    joblib.dump((predicted, actual), "pred_vs_actual.jbl")
                    torch.save(secondary_learner.state_dict(), igf_model_path)
                    best_test_loss = q_loss2

            secondary_learner.train()
    return secondary_learner

    Train the secondary learner (igf_model)

    Args:
        secondary_learner: secondary learner
        train_dataset: data to train secondary learner
        max_epochs: number of epochs to train secondary learner
        batch_size: batch size of training data of secondary learner
        eval_freq: secondary model evaluation can be triggered at eval_freq
        igf_model_path: path to store trained secondary learner

    Returns:
        Trained secondary learner

    
dpkd/transformers/examples/research_projects/information-gain-filtration/igf/igf.py