Fix multiprocessing dataloader checkpointing and use it in the train script

Fix multiprocessing dataloader checkpointing and use it in the train script #249