def compute_predictions(model, node_emb, seeds, device):
    """Compute the predictions for given source and destination nodes.

    This function computes the predictions for a set of node pairs, dividing the
    task into batches to handle potentially large graphs.
    """

    preds = torch.empty(seeds.shape[0], device=device)
    seeds_src, seeds_dst = seeds.T
    # The constant number is 1001, due to negtive ratio in the `ogbl-citation2`
    # dataset is 1000.
    eval_size = args.eval_batch_size * 1001
    # Loop over node pairs in batches.
    for start in trange(0, seeds_src.shape[0], eval_size, desc="Evaluate"):
        end = min(start + eval_size, seeds_src.shape[0])

        # Fetch embeddings for current batch of source and destination nodes.
        h_src = node_emb[seeds_src[start:end]].to(device, non_blocking=True)
        h_dst = node_emb[seeds_dst[start:end]].to(device, non_blocking=True)

        # Compute prediction scores using the model.
        preds[start:end] = predictions_step(model, h_src, h_dst)
    return preds
Compute the predictions for given source and destination nodes.

    This function computes the predictions for a set of node pairs, dividing the
    task into batches to handle potentially large graphs.
    
examples/graphbolt/pyg/link_prediction.py