Standard way to perform inference on LinkPrediction model

[johnloque]

Hi, I have successfully trained a link prediction model following this tutorial : Link Prediction on Heterogenous Graphs with PyG.ipynb. I have made a few changes (I've notably replaced the GNN part by a GAT), but the overall model structure and the use case remain the same.
However I've been struggling to understand how to perform inference for a given node (i.e. rank all possible links with respect to the classifier's output score).
I have found a workaround by defining a custom_sampler() function that takes a given node (i.e. user id) as input and returns a new HeteroData object containing its neighborhood (including all 1-hop and 2-hop neighbors). Then I'm setting the edge_index attribute with a tensor containing every possible link (i.e. a link for every movie I have in my dataset) to get classification scores that I can eventually use to retrieve the n most plausible links.
Although I couldn't find documentation about that, I guess there must be a more straightforward/standard/correct way to achieve such a thing using pyG objects ? Thanks for your help !

[rusty1s]

If you want to do prediction all possible links, I suggest you take a look at https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/recommender_system.py, which is an example especially designed for this us-ecase.