Use pin_memory and workers to speed up training

[Jorgvt]

Should try it out and check the time difference. Might be helpful before launching a ton of experiments.