Thanks for the quick reply, Shen Li. By using rpc_async, I got similar error:
$ mpirun -np 2 python rnn_jit.py
Traceback (most recent call last):
File “rnn_jit.py”, line 156, in 
run_worker()
File “rnn_jit.py”, line 145, in run_worker
_run_trainer()
File “rnn_jit.py”, line 116, in _run_trainer
script_model = torch.jit.script(model)
File “/opt/conda/lib/python3.7/site-packages/torch/jit/init.py”, line 1261, in script
return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 305, in create_script_module
return create_script_module_impl(nn_module, concrete_type, stubs_fn)
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 361, in create_script_module_impl
create_methods_from_stubs(concrete_type, stubs)
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 279, in create_methods_from_stubs
concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError:
rpc_async(dst_worker_name, user_callable, args, kwargs)does not support kwargs yet:
File “rnn_jit.py”, line 68
# pass input to the remote embedding table and fetch emb tensor back
# emb = _remote_method(EmbeddingTable.forward, self.emb_table_rref, input)
emb = rpc_async(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <— HERE
# output, hidden = self.rnn(emb, hidden)
# pass output to the rremote decoder and get the decoded output back
Line 68 refers to the rpc_async call inside the forward method
emb = rpc_async(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input)