How to do Link Prediction Inference

I have trained a hetero-graph successfully to do Link Prediction using EdgeDataLoader. Now, I’m trying to figure out whether a link should exist between 2 nodes or not (inference at test time), how can I do this? Further, let’s say I want to figure out given node A, what node is the most likely to connect to node A, and say even a step further and what are the 50 most likely nodes to connect to A. How can I do this?

I know how to get the node embeddings like the GraphSage example after the graph is trained.

Answered the same question in https://discuss.dgl.ai/t/running-link-prediction-on-disconnected-nodes-using-edgedataloader.