I was testing DataParallel in g5.12xlarge amazon linux AWS instance using this code. The code works perfectly if I just make one GPU visible, however it gets stuck when multiple GPUs are used.
System details:
torch version: 2.3.1
torch cuda version: 12.1
NVIDIA-SMI 555.42.06 Driver Version: 555.42.06 CUDA Version: 12.5