Feature: training checkpointing

Add checkpointing (and resuming from a checkpoint).