Code to reproduce the error:
import sys
import torch.distributed.rpc as rpc
import torch
import time
import torch.multiprocessing as mp
import torch.distributed.autograd as dist_autograd
import torch.nn as nn
from torch.distributed.rpc import RRef
from torch.distributed.optim import DistributedOptimizer
from torch import optim

def _call_method(method, rref, *args, **kwargs):
    r"""
    a helper function to call a method on the given RRef
    """
    return method(rref.local_value(), *args, **kwargs)


def _remote_method(method, rref, *args, **kwargs):
    r"""
    a helper function to run method on the owner of rref and fetch back the
    result using RPC
    """
    return rpc.rpc_sync(
        rref.owner(),
        _call_method,
        args=[method, rref] + list(args),
        kwargs=kwargs
    )

class Net1(nn.Module):
    def __init__(self):
        super(Net1, self).__init__()
        self.layer = nn.Linear(10, 20)

    def parameter_rrefs(self):
        return [RRef(p) for p in self.parameters() if p.requires_grad]

    def forward(self, x):
        return self.layer(x)

class Net2(nn.Module):
    def __init__(self):
        super(Net2, self).__init__()
        self.layer = nn.Linear(20, 1)

    def parameter_rrefs(self):
        return [RRef(p) for p in self.parameters() if p.requires_grad]

    def forward(self, x):
        return self.layer(x)

class Net(nn.Module):
    def __init__(self, *args, **kwargs):
        super(Net, self).__init__()
        self.encoder_rref = rpc.remote(
            "worker1",
            Net1,
            args = args,
            kwargs = kwargs
        )

        self.decoder_rref = rpc.remote(
            "worker2",
            Net2,
            args = args,
            kwargs = kwargs
        )

    def parameter_rrefs(self):
        remote_params = []
        remote_params.extend(self.encoder_rref.remote().parameter_rrefs().to_here())
        remote_params.extend(self.decoder_rref.remote().parameter_rrefs().to_here())
        return remote_params

    def forward(self, x):
        x = _remote_method(Net1.forward, self.encoder_rref, x)
        x = _remote_method(Net2.forward, self.decoder_rref, x)
        return x

def run_master():
    model = Net()
    opt = DistributedOptimizer(
        optim.SGD,
        model.parameter_rrefs(),
        lr=0.05,
    )
    for i in range(10):
        with dist_autograd.context() as context_id:
            x = torch.randn(32, 10)
            loss = model(x)
            loss = loss.sum()
            print('Before dist_autograd')
            dist_autograd.backward(context_id, [loss])
            opt.step(context_id)

def workers_init_rpc(rank, world_size, options):
    print(f'Rank {rank}: Proceed to init rpc')
    rpc.init_rpc(
        f"worker{rank}",
        rank=rank,
        world_size=world_size,
        rpc_backend_options=options
    )
    print(f'Rank: {rank}, rpc init done')

    if rank == 0:
        print('Proceed to run_master')
        run_master()
    rpc.shutdown()


if __name__=="__main__":
    world_size = 3
    processes = []

    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=16, rpc_timeout = 60*60)
    rank = int(sys.argv[1])

    p = mp.Process(target=workers_init_rpc, args=(rank, world_size, options))

    p.start()

    processes.append(p)

    for p in processes:
        p.join()

Save this in a .py file.
How to run:
Env. variable exports for all of these 3 nodes:
export MASTER_ADDR=<Node 0 IP>
export MASTER_PORT=<Node 0 port>
export GLOO_SOCKET_IFNAME=network interface
export TP_SOCKET_IFNAME=network interface
On Node 0: python <filename> 0
On Node 1: python <filename> 1
On Node 2: python <filename> 2
PyTorch Version: 1.7.1