To extract the embedding of individual atoms

[nonsensejoke]

Hi there, thank you for sharing this exciting work!

I wonder how can I obtain the MolE atomic embeddings for individual atoms as mentioned in the Figure 2b of the paper?

Thanks in advance.