Solved by Chenguang_Wan in post #3 
              
              
                I find the solution by the help with @ngimel. This is the solution learned from her. 

def demo_basic(rank, world_size):

    print(f"Running basic DDP example on rank {rank}.")

    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    torch.cuda.set_device(rank)

    ToyModel = nn…