def METHOD_NAME(model, g, A_hat, X):
    label = g.ndata["label"]
    train_mask = g.ndata["train_mask"]
    optimizer = Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
    for epoch in range(50):
        # Forward.
        model.METHOD_NAME()
        logits = model(A_hat, X)
        # Compute loss with nodes in training set.
        loss = F.cross_entropy(logits[train_mask], label[train_mask])
        # Backward.
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # Compute prediction.
        model.eval()
        logits = model(A_hat, X)
        pred = logits.argmax(dim=1)
        # Evaluate the prediction.
        val_acc, test_acc = evaluate(g, pred)
        print(
            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
            f" acc: {test_acc:.3f}"
        )