with dist_autograd.context() as context_id:
           # Forward pass (create references on remote nodes).
           rref1 = rpc.remote(dst_name, random_tensor)
           rref2 = rpc.remote(dst_name, random_tensor)
           loss = rref1.to_here() + rref2.to_here()
       # Backward pass (run distributed autograd).
       dist_autograd.backward(context_id, [loss.sum()])

       # Build DistributedOptimizer.
       dist_optim = DistributedOptimizer(
       optim.SGD,
       [rref1, rref2],
       lr=0.05,
       )

       # Run the distributed optimizer step.
       dist_optim.step(context_id)

def run_process(rank, world_size):
       dst_rank = (rank + 1) % world_size
       _run_process(rank, dst_rank, world_size)
       rpc.shutdown()
if name == 'main':
     # Run world_size workers
     world_size = 2
     mp.spawn(run_process, args=(world_size,), nprocs=world_size)