Cleanup

[thomas chaton]

Dear @MilesCranmer,

Awesome work with this repo.

I really like how you are using jax to generate the data and PyTorch Geometric for the model.
I think you could have made your code even cleaner if you used PyTorch Lightning.
We worked with Matthias Fey to better support PyTorch Geometric and working DGL team right now for the same purpose.

Find examples there: https://github.com/rusty1s/pytorch_geometric/blob/master/examples/pytorch_lightning/graph_sage.py#L72

Best,
T.C