Torch_geometric node2vec get vector

For example if I have a graph of 3 nodes A, B, and C. What I want to do is to feed this graph into a node2vec model and learn an embedding of each node. Here is the pipe line

  1. Construct a graph using nextworkx
  2. Convert it to a pytorch graph data using from_networkx
  3. Construct a node2vec model and train it.
    After training I should get a 3XN matrix where N is the embedding dimension I desired. My question is, how do I know which of the 3 embedding is the embedding for node A