[FEATURE REQUEST] Rewrite training code based on PyTorch Lightning

[Jn-Huang]

Is your feature request related to a problem? Please describe.
Using lightning can simplify our implementation of models, training, etc. The idea is from @xingjian-zhang. I am still going through the tutorials of lightning.

We may end up with a trainer class similar to this one:
https://github.com/intel-analytics/BigDL/blob/main/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py

[xingjian-zhang]

The Trainer class is a bit complicated to me. Do we really need to implement a trainer here? My initial idea was, for each kind of task, to implement a subclass of pl.LightningModule that at least reloads/implements the following member functions:

1. forward
2. training_step
3. validation_step
4. test_step
5. configure_optimizers
   We may want to add other functions by need, too.

Therefore, the training script will contain these components:

1. A subclass of pl.LightningModule
2. Class of GNN(s) (either imported or implemented)
3. A main function that
   1. Instantiates 1 and 2
   2. Gets the dataloader
   3. Instantiates a pl.Trainer to fit train data
   4. Reports results on test data

See details at https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html