Improvements to distributed training

There was a bug in distributed training when using more than one GPU causing training to stall at the end of the last epoch.