The question seems a little unclear but essentially I think your problem has a solution that can be reduced to the below points.

1. Graph-Level Prediction Needed: Current setup makes node-level predictions, but task requires a single prediction per graph.

2. Loss Calculation Misalignment: BCEWithLogitsLoss applied to all nodes; should focus on the "correct" node or aggregate node embeddings.

3. Label-Output Mismatch: Ensure truth tensor in train() and val() matches intended output for only the "correct" node

4. Rank Calculation Issue in test(): Only rank the "correct" node or relevant nodes, not all nodes.

5. Pooling Layer for Aggregation: Use a global pooling layer (e.g., global_mean_pool) to create a graph-level embedding.