Support gradient checkpointing

#3
by maxall4 - opened