[rpc][jit] support remote call in TorchScript (#43046)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/43046

Test Plan: Imported from OSS

Reviewed By: mrshenli

Differential Revision: D23621108

Pulled By: wanchaol

fbshipit-source-id: e8152c6cdd3831f32d72d46ac86ce22f3f13c651
This commit is contained in:
Wanchao Liang
2020-09-11 14:48:45 -07:00
committed by Facebook GitHub Bot
parent 3e5df5f216
commit ab6126b50e
9 changed files with 49 additions and 7 deletions

View File

@ -864,6 +864,11 @@ std::shared_ptr<SugaredValue> toSugaredValue(
obj.ptr() ==
py::module::import("torch.distributed.rpc").attr("rpc_sync").ptr()) {
return SpecialFormValue::create(prim::rpc_sync);
} else if (
// RPC module is only avaialble when build flag "USE_DISTRIBUTED" is on.
obj.ptr() ==
py::module::import("torch.distributed.rpc").attr("remote").ptr()) {
return SpecialFormValue::create(prim::rpc_remote);
#endif
} else if (auto callee = as_module(obj)) {
throw ErrorReport(loc) << "Cannot call a ScriptModule that is not"