Fix a bug in gene_random_walk, import pyg only when initing Raindrop, and make MultiHeadAttention work with all attention operators

[WenjieDu]

What does this PR do?

1. fixing a bug in gene_random_walk that train_X_ori and val_X_ori didn't get standardized;
2. fixing Only reminder missing torch_geometric when initializing Raindrop #380;
3. fixing Make MultiHeadAttention to work with all attention operators #382;

Before submitting

• This PR is made to fix a typo or improve the docs (you can dismiss the other checks if this is the case).
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• I have commented my code, particularly in hard-to-understand areas.
• I have written necessary tests and already run them locally.