def METHOD_NAME(params):
    dataset = MolHIVDataset()
    accelerator.print(
        f"train, test, val sizes: {len(dataset.train)}, "
        f"{len(dataset.test)}, {len(dataset.val)}."
    )
    accelerator.print("Finished loading.")
    train_loader = GraphDataLoader(
        dataset.train,
        batch_size=params.batch_size,
        shuffle=True,
        collate_fn=dataset.collate,
        pin_memory=True,
        num_workers=16,
    )
    val_loader = GraphDataLoader(
        dataset.val,
        batch_size=params.batch_size,
        shuffle=False,
        collate_fn=dataset.collate,
        pin_memory=True,
        num_workers=16,
    )
    test_loader = GraphDataLoader(
        dataset.test,
        batch_size=params.batch_size,
        shuffle=False,
        collate_fn=dataset.collate,
        pin_memory=True,
        num_workers=16,
    )
    # Load pre-trained model.
    download(url="https://data.dgl.ai/pre_trained/graphormer_pcqm.pth")
    model = Graphormer()
    state_dict = th.load("graphormer_pcqm.pth")
    model.load_state_dict(state_dict)
    model.reset_output_layer_parameters()
    num_epochs = 16
    total_updates = 33000 * num_epochs / params.batch_size
    # Use warmup schedule to avoid overfitting at the very beginning
    # of training, the ratio 0.16 is the same as the paper.
    warmup_updates = total_updates * 0.16
    optimizer = AdamW(model.parameters(), lr=1e-4, eps=1e-8, weight_decay=0)
    lr_scheduler = get_polynomial_decay_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_updates,
        num_training_steps=total_updates,
        lr_end=1e-9,
        power=1.0,
    )
    epoch_train_AUCs, epoch_val_AUCs, epoch_test_AUCs = [], [], []
    # Pass all objects relevant to training to the prepare() method as required
    # by Accelerate.
    (
        model,
        optimizer,
        train_loader,
        val_loader,
        test_loader,
        lr_scheduler,
    ) = accelerator.prepare(
        model, optimizer, train_loader, val_loader, test_loader, lr_scheduler
    )
    for epoch in range(num_epochs):
        epoch_train_loss, epoch_train_auc = train_epoch(
            model, optimizer, train_loader, lr_scheduler
        )
        epoch_val_loss, epoch_val_auc = evaluate_network(model, val_loader)
        epoch_test_loss, epoch_test_auc = evaluate_network(model, test_loader)
        epoch_train_AUCs.append(epoch_train_auc)
        epoch_val_AUCs.append(epoch_val_auc)
        epoch_test_AUCs.append(epoch_test_auc)
        accelerator.print(
            f"Epoch={epoch + 1} | train_AUC={epoch_train_auc:.3f} | "
            f"val_AUC={epoch_val_auc:.3f} | test_AUC={epoch_test_auc:.3f}"
        )
    # Return test and train AUCs with best val AUC.
    index = epoch_val_AUCs.index(max(epoch_val_AUCs))
    val_auc = epoch_val_AUCs[index]
    train_auc = epoch_train_AUCs[index]
    test_auc = epoch_test_AUCs[index]
    accelerator.print("Test ROCAUC: {:.4f}".format(test_auc))
    accelerator.print("Val ROCAUC: {:.4f}".format(val_auc))
    accelerator.print("Train ROCAUC: {:.4f}".format(train_auc))
    accelerator.print("Best epoch index: {:.4f}".format(index))