Distributed training hangs