This script is a repro for the issue
import torch
import os
import torch.distributed as dist
from torch.distributed.rpc import rpc_sync
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed.rpc as rpc
from torch import optim #XXX include this
from torch.distributed.optim import DistributedOptimizer #XXX include this
from torch.distributed.rpc import RRef, rpc_async, remote
import torch.distributed.autograd as dist_autograd #XXX include this

def get_local_rank():
    return int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])

def _parameter_rrefs(module):
    param_rrefs = []
    for param in module.parameters():
        param_rrefs.append(RRef(param))
    return param_rrefs

class EmbeddingTable(nn.Module):
    r"""
    Encoding layers of the RNNModel
    """
    def __init__(self, ntoken, ninp, dropout):
        super(EmbeddingTable, self).__init__()
        self.drop = nn.Dropout(dropout)
        self.encoder = nn.Embedding(ntoken, ninp).cuda()
        self.encoder.weight.data.uniform_(-0.1, 0.1)

    def forward(self, input):
        return self.drop(self.encoder(input.cuda())).cpu() # XXX: extra ')'


class Decoder(nn.Module):
    def __init__(self, ntoken, nhid, dropout):
        super(Decoder, self).__init__()
        self.drop = nn.Dropout(dropout)
        self.decoder = nn.Linear(nhid, ntoken)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-0.1, 0.1)

    def forward(self, output):
        return self.decoder(self.drop(output))

class RNNModel(nn.Module):
    def __init__(self, ps, ntoken, ninp, nhid, nlayers, dropout=0.5):
        super(RNNModel, self).__init__()

        # setup embedding table remotely
        self.emb_table_rref = rpc.remote(ps, EmbeddingTable, args=(ntoken, ninp, dropout))
        # setup LSTM locally
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
        # setup decoder remotely
        self.decoder_rref = rpc.remote(ps, Decoder, args=(ntoken, nhid, dropout))

    def forward(self, input, hidden):
        # pass input to the remote embedding table and fetch emb tensor back
        # emb = _remote_method(EmbeddingTable.forward, self.emb_table_rref, input) # Original call
        emb = rpc_async(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input) # adapted call

        output, hidden = self.rnn(emb, hidden)

        # pass output to the rremote decoder and get the decoded output back
        # decoded = _remote_method(Decoder.forward, self.decoder_rref, output) # Original call
        decoded = rpc_async(self.decoder_rref.owner(), Decoder.forward, args=output) # adapted call
        return decoded, hidden
    def parameter_rrefs(self):
        remote_params = []
        # get RRefs of embedding table
        # remote_params.extend(_remote_method(_parameter_rrefs, self.emb_table_rref)) # Original call
        remote_params.extend(rpc_async(_parameter_rrefs, self.emb_table_rref))

        # create RRefs for local parameters
        remote_params.extend(_parameter_rrefs(self.rnn))

        # get RRefs of decoder
        # remote_params.extend(_remote_method(_parameter_rrefs, self.decoder_rref)) # Original call
        remote_params.extend(rpc_async(_parameter_rrefs, self.decoder_rref)) # Adapted call
        return remote_params

def _run_trainer():
    ntoken = 10
    ninp = 2
    nhid = 3
    nlayers = 4
    model = RNNModel('ps', ntoken, ninp, nhid, nlayers) # XXX: no rnn.
    script_model = torch.jit.script(model)
    print(script_model)

def run_worker():
    world_size=2
    rank=get_local_rank()
    os.environ['MASTER_ADDR'] = '10.123.134.28'
    os.environ['MASTER_PORT'] = '21234'
    if rank == 1:
        rpc.init_rpc("trainer", rank=rank, world_size=world_size)
        _run_trainer()
    else:
        rpc.init_rpc("ps", rank=rank, world_size=world_size)
        # parameter server do nothing
        pass
    # block until all rpcs finish
    rpc.shutdown()

if __name__=="__main__":
    run_worker()