[RPC][BUGFIX] Fix remote device sync (#5538)

diff --git a/src/runtime/rpc/rpc_device_api.cc b/src/runtime/rpc/rpc_device_api.cc
index ade4d16..93af4e2 100644 (file)
--- 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 (file)
--- 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)