Hello @ptrblck  thanks for your answer,
I tried using the dockerfile here . I also tried to debug whether Ubuntu version was the reason, which isn’t as the code gets stuck on 18.04 as well. Can you please suggest what else I could try to fix this? Attached below are terminal outputs,
$ docker run --rm -it --init --gpus=2 --ipc=host --volume="$PWD:/app" anibali/pytorch:1.8.1-cuda11.1 python3 -m torch.distributed.launch --nproc_per_node=2 gpu_test_script.py
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
fc1.weight Parameter containing:
tensor([[ 0.2373, -0.6326],
        [-0.5230,  0.0513]], device='cuda:0', requires_grad=True)
x tensor([[-0.4446,  0.0000]], device='cuda:0', grad_fn=<FusedDropoutBackward>)
iter 0, out tensor([[-0.4446,  0.0000]], device='cuda:0', grad_fn=<FusedDropoutBackward>)
^CKilling subprocess 9
Killing subprocess 10
Main process received SIGINT, exiting

$ docker run --rm -it --init --gpus=2 --ipc=host --volume="$PWD:/app" cuda11_ubuntu1804 python3 -m torch.distributed.launch --nproc_per_node=2 gpu_test_script.py
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
fc1.weight Parameter containing:
tensor([[-0.1156, -0.3040],
        [-0.5875,  0.4592]], device='cuda:0', requires_grad=True)
x tensor([[0., 0.]], device='cuda:0', grad_fn=<FusedDropoutBackward>)
iter 0, out tensor([[0., 0.]], device='cuda:0', grad_fn=<FusedDropoutBackward>)

Thank you again for your time patience and help @ptrblck , much appreciate it.