Add metrics to collect during training

[wanchaol]

see https://github.com/pytorch-labs/torchtrain/blob/main/train.py#L87

we should have the following metrics associated with the train steps:

1. gpu memory usage
2. wps
3. loss

[lessw2020]

I can work on this one, with gpu stats first.
I think we also want mfu.

[lessw2020]

and also flop counter

[gnadathur]

@lessw2020 , @tianyu-l -- Can we close this issue ?

[gnadathur]

closing as per conversation w/ @lessw2020