Hi, Rohan
Thanks for your attention. I’ve managed to resolve the hang by adding torch.cuda.set_device(rank) before the training loop. This stopped subprocesses with rank larger than 0 from allocating memory on cuda:0 i.e. the device used for subprocess with rank 0.
Cheers,
Fred