From b4ee81d92849ec41e97e137ff1853d8aadd51ab1 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 7 May 2020 14:54:25 -0700 Subject: [PATCH] [RPC][BUGFIX] Fix remote device sync (#5538) --- src/runtime/rpc/rpc_device_api.cc | 2 +- tests/python/unittest/test_runtime_rpc.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/runtime/rpc/rpc_device_api.cc b/src/runtime/rpc/rpc_device_api.cc index ade4d16..93af4e2 100644 --- a/src/runtime/rpc/rpc_device_api.cc +++ b/src/runtime/rpc/rpc_device_api.cc @@ -109,7 +109,7 @@ class RPCDeviceAPI final : public DeviceAPI { void StreamSync(TVMContext ctx, TVMStreamHandle stream) final { auto remote_ctx = RemoveSessMask(ctx); - GetSess(ctx)->GetDeviceAPI(remote_ctx)->StreamSync(ctx, stream); + GetSess(ctx)->GetDeviceAPI(remote_ctx)->StreamSync(remote_ctx, stream); } private: diff --git a/tests/python/unittest/test_runtime_rpc.py b/tests/python/unittest/test_runtime_rpc.py index 58cd8e2..6a46eab 100644 --- a/tests/python/unittest/test_runtime_rpc.py +++ b/tests/python/unittest/test_runtime_rpc.py @@ -128,6 +128,8 @@ def test_rpc_echo(): "testing.test_raise_error_callback")("RuntimeError") raise_err() + remote.cpu().sync() + temp = rpc.server._server_env([]) server = rpc.Server("localhost") client = rpc.connect(server.host, server.port) -- 2.7.4