Dataparallel gets stuck in AWS

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