frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rpc][jit] support rpc_sync in TorchScript (#43043)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43043

This add the support for rpc_sync in TorchScript in a way similar to
rpc_async

Test Plan: Imported from OSS

Reviewed By: mrshenli

Differential Revision: D23252039

Pulled By: wanchaol

fbshipit-source-id: 8a05329cb8a24079b2863178b73087d47273914c
This commit is contained in:
Wanchao Liang
2020-09-11 14:48:45 -07:00
committed by Facebook GitHub Bot
parent 8bec7cfa91
commit 3e5df5f216
9 changed files with 220 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/csrc/jit/python/python_sugared_value.cpp
View File

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