Checkpoint options

Checkpoint update: https://pytorch.org/docs/stable/checkpoint.html. It will default the use_reentrant to False. We should check the impact of this use_reentrant option on 1) training time, 2) model performance, 3) memory usage.