From 521cf8b9f440f20dc26bba6a3589ba18958329eb Mon Sep 17 00:00:00 2001 From: Jampani Date: Tue, 15 Aug 2017 18:32:31 -0700 Subject: [PATCH] Fixes scaling issue identified in #7455 --- src/kvstore/kvstore_dist.h | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/kvstore/kvstore_dist.h b/src/kvstore/kvstore_dist.h index 52c7c132cb5c..e313408732ce 100644 --- a/src/kvstore/kvstore_dist.h +++ b/src/kvstore/kvstore_dist.h @@ -232,17 +232,15 @@ class KVStoreDist : public KVStoreLocal { } // push to servers - send_buf.WaitToRead(); - size_t size = send_buf.shape().Size(); + auto push_to_servers = + [this, key, &send_buf](RunContext rctx, Engine::CallbackOnComplete cb) { + size_t size = send_buf.shape().Size(); #if MKL_EXPERIMENTAL == 1 - mkl_set_tblob_eager_mode(send_buf.data()); + mkl_set_tblob_eager_mode(send_buf.data()); #endif - real_t* data = static_cast(send_buf.data().dptr_); - auto push_to_servers = - [this, key, data, size](RunContext rctx, Engine::CallbackOnComplete cb) { + real_t* data = static_cast(send_buf.data().dptr_); // convert to ps keys PSKV& pskv = EncodeKey(key, size); - // do push. false means no delete ps::SArray vals(data, size, false); CHECK_NOTNULL(ps_worker_)->ZPush(