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

trevor-m · Jun 8, 2020 · 476623a · 476623a
1 parent 132017d
commit 476623a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git 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
@@ -128,6 +128,8 @@ def check(remote):
                 "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)